import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import os
from huggingface_hub import login

hf_token = os.getenv("HF_TOKEN")  # Get token from environment variable
login(hf_token)



# Load your model and tokenizer
model_name = "Amir230703/phi3-medmcqa-finetuned"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")

# Define function for inference
def generate_response(prompt):
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
    with torch.no_grad():
        output = model.generate(input_ids, max_length=200)
    return tokenizer.decode(output[0], skip_special_tokens=True)

# Create Gradio UI
iface = gr.Interface(
    fn=generate_response,
    inputs="text",
    outputs="text",
    title="Medical QA Model",
    description="Enter a medical question, and the AI will provide an answer."
)

iface.launch()