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

# Retrieve the token from environment variables
api_token = os.getenv("HF_TOKEN")

# Load the Hugging Face model and tokenizer with authentication
model_name = "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1"
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=api_token)
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, use_auth_token=api_token)


# Define the function to process user input
def generate_response(input_text):
    try:
        # Tokenize the input text
        inputs = tokenizer(input_text, return_tensors="pt")
        
        # Generate a response using the model
        outputs = model.generate(
            inputs["input_ids"],
            max_length=256,  # Limit the output length
            num_return_sequences=1,  # Generate a single response
            temperature=0.7,  # Adjust for creativity vs. determinism
            top_p=0.9,  # Nucleus sampling
            top_k=50  # Top-k sampling
        )
        
        # Decode and return the generated text
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return response
    
    except Exception as e:
        return f"Error: {str(e)}"

# Create a Gradio interface with API enabled
iface = gr.Interface(
    fn=generate_response,
    inputs="text",
    outputs="text",
    title="ContactDoctor Medical Assistant",
    description="Provide input symptoms or queries and get AI-powered medical advice.",
    enable_api=True  # Enables API for external calls
)

# Launch the Gradio app
if __name__ == "__main__":
    iface.launch()