| import gradio as gr | |
| from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline | |
| import torch | |
| model_id = "aaditya/Llama3-OpenBioLLM-8B" | |
| tokenizer = AutoTokenizer.from_pretrained(model_id) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| model_id, | |
| device_map="auto", | |
| torch_dtype=torch.float16 | |
| ) | |
| chat_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer) | |
| def chatbot(message, history=[]): | |
| prompt = f"[INST] {message} [/INST]" | |
| response = chat_pipeline(prompt, max_new_tokens=256, do_sample=True, temperature=0.7)[0]['generated_text'] | |
| return response.replace(prompt, "").strip() | |
| gr.ChatInterface(fn=chatbot, title="🩺 OpenBioLLM Chatbot", description="Ask me anything biomedical!").launch() | |