import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Load the model and tokenizer
model_name = "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Add pad token if it doesn't exist
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

def skin_disease_assistant(user_input):
    # Create a medical-focused prompt
    prompt = f"""You are a medical AI assistant specializing in skin diseases. Provide helpful, accurate information about skin conditions based on the following query.

Query: {user_input}

Response:"""
    
    # Tokenize and generate response
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
    
    with torch.no_grad():
        outputs = model.generate(
            inputs.input_ids,
            max_new_tokens=200,
            temperature=0.7,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id
        )
    
    # Decode the response
    full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # Extract only the generated part (after "Response:")
    response = full_response.split("Response:")[-1].strip()
    
    return response

# Create Gradio interface
demo = gr.Interface(
    fn=skin_disease_assistant,
    inputs=gr.Textbox(label="Ask about skin conditions", placeholder="e.g., What are the symptoms of eczema?"),
    outputs=gr.Textbox(label="AI Response"),
    title="🏥 Skin Disease AI Assistant",
    description="An AI assistant to help with skin disease questions. For educational purposes only - consult medical professionals for actual diagnosis.",
    examples=[
        "What are the symptoms of psoriasis?",
        "How to treat acne?",
        "What causes eczema flare-ups?",
        "Difference between melanoma and mole?"
    ]
)

demo.launch()