docto/Docto-Bot

1import torch2import gradio as gr3from transformers import AutoTokenizer, AutoModelForCausalLM4 5MODEL = "docto/Docto-Bot"6 7tokenizer = AutoTokenizer.from_pretrained(MODEL)8model = AutoModelForCausalLM.from_pretrained(MODEL)9 10device = "cuda" if torch.cuda.is_available() else "cpu"11model.to(device)12 13if tokenizer.pad_token is None:14    tokenizer.pad_token = tokenizer.eos_token15 16 17def get_reply(user_input):18 19    prompt = f"Question: {user_input}\nAnswer:"20 21    inputs = tokenizer(22        prompt,23        return_tensors="pt"24    ).to(device)25 26    outputs = model.generate(27        **inputs,28        max_new_tokens=150,29        do_sample=True,30        temperature=0.7,31        top_k=50,32        top_p=0.9,33        repetition_penalty=1.15,34        no_repeat_ngram_size=3,35        pad_token_id=tokenizer.eos_token_id,36        eos_token_id=tokenizer.eos_token_id37    )38 39    response = tokenizer.decode(40        outputs[0],41        skip_special_tokens=True42    )43 44    if "Answer:" in response:45        response = response.split("Answer:", 1)[1]46 47    return response.strip()48 49 50iface = gr.Interface(51    fn=get_reply,52 53    inputs=gr.Textbox(54        lines=2,55        placeholder="Ask a medical question..."56    ),57 58    outputs=gr.Textbox(59        label="Response"60    ),61 62    title="Docto-Bot",63 64    description="Medical Question Answering Bot"65)66 67iface.launch()