| import torch |
| import gradio as gr |
| from transformers import AutoTokenizer, AutoModelForCausalLM |
|
|
| MODEL = "docto/Docto-Bot" |
|
|
| tokenizer = AutoTokenizer.from_pretrained(MODEL) |
| model = AutoModelForCausalLM.from_pretrained(MODEL) |
|
|
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| model.to(device) |
|
|
| if tokenizer.pad_token is None: |
| tokenizer.pad_token = tokenizer.eos_token |
|
|
|
|
| def get_reply(user_input): |
|
|
| prompt = f"Question: {user_input}\nAnswer:" |
|
|
| inputs = tokenizer( |
| prompt, |
| return_tensors="pt" |
| ).to(device) |
|
|
| outputs = model.generate( |
| **inputs, |
| max_new_tokens=150, |
| do_sample=True, |
| temperature=0.7, |
| top_k=50, |
| top_p=0.9, |
| repetition_penalty=1.15, |
| no_repeat_ngram_size=3, |
| pad_token_id=tokenizer.eos_token_id, |
| eos_token_id=tokenizer.eos_token_id |
| ) |
|
|
| response = tokenizer.decode( |
| outputs[0], |
| skip_special_tokens=True |
| ) |
|
|
| if "Answer:" in response: |
| response = response.split("Answer:", 1)[1] |
|
|
| return response.strip() |
|
|
|
|
| iface = gr.Interface( |
| fn=get_reply, |
|
|
| inputs=gr.Textbox( |
| lines=2, |
| placeholder="Ask a medical question..." |
| ), |
|
|
| outputs=gr.Textbox( |
| label="Response" |
| ), |
|
|
| title="Docto-Bot", |
|
|
| description="Medical Question Answering Bot" |
| ) |
|
|
| iface.launch() |
