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").strip()

# Model name
model_name = "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1"

# Load the Hugging Face model and tokenizer with required arguments
tokenizer = AutoTokenizer.from_pretrained(
    model_name,
    token=api_token,  # Authenticate with Hugging Face token
    trust_remote_code=True  # Allow custom code from the repository
)

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    token=api_token,
    trust_remote_code=True,
    device_map="auto",  # Efficient device allocation
    torch_dtype=torch.float16  # Mixed precision for faster inference
)

# 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()