app.py

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import os
print("HUGGINGFACEHUB_API_TOKEN:", os.environ.get("HUGGINGFACEHUB_API_TOKEN"))

MODEL_NAME = "meta-llama/Meta-Llama-3-8B-Instruct"

SYSTEM_PROMPT = "You are a helpful doctor. Answer the user's question."

# 3-6 diverse medical Q&A style examples
EXAMPLES = [
    {
        "input": "I have a sore throat and mild fever.",
        "output": "These symptoms could be due to a common viral infection. Rest, drink fluids, and monitor your symptoms. If you develop high fever or difficulty swallowing, consult a healthcare provider."
    },
    {
        "input": "I'm experiencing chest pain after eating.",
        "output": "Chest pain after eating may be due to acid reflux or indigestion. Try to avoid heavy meals and spicy foods. If the pain is severe or associated with shortness of breath, seek medical attention."
    },
    {
        "input": "My child has a rash and a high temperature.",
        "output": "A rash with a high temperature in a child could be due to a viral illness, but some rashes can indicate serious conditions. Monitor for other symptoms and consult a doctor promptly."
    },
    {
        "input": "I feel very tired all the time.",
        "output": "Persistent tiredness can be caused by stress, poor sleep, or an underlying medical problem. Make sure you're resting, eating well, and consider seeing a doctor if it continues."
    },
    {
        "input": "I'm having trouble sleeping at night.",
        "output": "Trouble sleeping can be due to stress or changes in routine. Try to maintain a regular sleep schedule and avoid caffeine in the evening."
    },
    {
        "input": "I have a headache with sensitivity to light.",
        "output": "A headache with sensitivity to light may be a migraine. Rest in a dark, quiet room and consider over-the-counter pain relief. If the headache is severe or sudden, seek medical care."
    }
]

def build_prompt(user_input):
    prompt = (
        "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n"
        f"{SYSTEM_PROMPT}\n"
    )
    for ex in EXAMPLES:
        prompt += (
            "<|start_header_id|>user<|end_header_id|>\n"
            f"Patient: {ex['input']}\n"
            "<|start_header_id|>assistant<|end_header_id|>\n"
            f"Doctor: {ex['output']}\n\n"
        )
    prompt += (
        "<|start_header_id|>user<|end_header_id|>\n"
        f"Patient: {user_input}\n"
        "<|start_header_id|>assistant<|end_header_id|>\n"
        "Doctor:"
    )
    return prompt

# Load model directly
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")

model.eval()

def chat(input_text):
    prompt = build_prompt(input_text)
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
    with torch.no_grad():
        output = model.generate(
            input_ids,
            max_new_tokens=128,
            do_sample=True,
            temperature=0.7,
            top_p=0.95,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    response = tokenizer.decode(output[0][input_ids.shape[-1]:], skip_special_tokens=True)
    # Keep only up to the first double linebreak or sentence for cleanliness
    response = response.strip().split("\n")[0]
    if "." in response:
        response = response.split(".")[0] + "."
    return response

demo = gr.Interface(
    fn=chat,
    inputs=gr.Textbox(lines=2, placeholder="Enter your medical question..."),
    outputs="text",
    title="Llama-3-8B Medical Q&A",
    description="Ask a medical question to the Llama-3-8B model! Example: 'I have bloating with chest pain.'"
)

if __name__ == "__main__":
    demo.launch()