First commit

app.py

@@ -0,0 +1,69 @@
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+import torch
+model_name = "ruslanmv/Medical-Llama3-8B"
+device_map = 'auto'
+# Check if GPU is available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
+if device.type == "cuda":
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.float16,
+    )
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        quantization_config=bnb_config,
+        trust_remote_code=True,
+        use_cache=False,
+        device_map=device_map
+    )
+else:
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        trust_remote_code=True,
+        use_cache=False
+    )
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+tokenizer.pad_token = tokenizer.eos_token
+
+def askme(symptoms, question):
+    sys_message = '''
+    You are an AI Medical Assistant trained on a vast dataset of health information. Please be thorough and
+    provide an informative answer. If you don't know the answer to a specific medical inquiry, advise seeking professional help.
+    '''
+    content = symptoms + " " + question
+    messages = [{"role": "system", "content": sys_message}, {"role": "user", "content": content}]
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    outputs = model.generate(**inputs, max_new_tokens=200, use_cache=True)
+    response_text = tokenizer.batch_decode(outputs)[0].strip()
+    answer = response_text.split('<|im_start|>assistant')[-1].strip()
+    return answer
+
+# Example usage
+symptoms = '''
+I'm a 35-year-old male and for the past few months, I've been experiencing fatigue,
+increased sensitivity to cold, and dry, itchy skin.
+'''
+question = '''
+Could these symptoms be related to hypothyroidism?
+If so, what steps should I take to get a proper diagnosis and discuss treatment options?
+'''
+
+examples = [
+    [symptoms, question]
+]
+
+iface = gr.Interface(
+    fn=askme,
+    inputs=["text", "text"],
+    outputs="text",
+    examples=examples,
+    title="Medical AI Chatbot",
+    description="Ask me a medical question!"
+)
+
+iface.launch()

requirements.txt

@@ -0,0 +1,8 @@
+torch==2.2.1
+torchvision
+torchaudio
+xformers
+bitsandbytes
+accelerate
+gradio
+transformers