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

# Load pretrained model and tokenizer
model_name = "jrocha/tiny_llama"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Load data
df = pd.read_csv('splitted_df_jo.csv')

# Function to prepare context
def prepare_context():
    pubmed_information_column = df['section_text']
    pubmed_information_cleaned = ""
    for text in pubmed_information_column.tolist():
        objective_index = text.find("Objective")
        if objective_index != -1:
            cleaned_text = text[:objective_index]
            pubmed_information_cleaned += cleaned_text
        else:
            pubmed_information_cleaned += text
    max_length = 1000  # Adjust as needed
    return pubmed_information_cleaned[:max_length]

# Function to generate answer
def answer_question(question):
    pubmed_information_cleaned = prepare_context()

    # Prepare input sequence
    messages = [
        {
            "role": "system",
            "content": "You are a friendly chatbot who responds to questions about cancer. Please be considerate.",
        },
        {"role": "user", "content": question},
    ]
    prompt_with_pubmed = f"{pubmed_information_cleaned}\n\n"  # Adjust formatting as needed
    prompt_with_pubmed += tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)

    # Generate response
    input_ids = tokenizer.encode(prompt_with_pubmed, return_tensors='pt')
    output = model.generate(input_ids, max_length=600, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)

    # Decode and return generated text
    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
    position_assistant = generated_text.find("<|assistant|>") + len("<|assistant|>")
    return generated_text[position_assistant:]

def main():
    """"
    Initializes a Women Cancer ChatBot interface using Hugging Face models for question answering.
    
    This function loads a pretrained tokenizer and model from the Hugging Face model hub
    and creates a Gradio interface for the ChatBot. Users can input questions related to
    women's cancer topics, and the ChatBot will generate answers based on the provided context.
    
    Returns:
    None
    Example:
    >>> main()
    """
    iface = gr.Interface(fn=answer_question,
                         inputs=["text"],
                         outputs=[gr.Textbox(label="Answer")],
                         title="Women Cancer ChatBot",
                         description="How can I help you?",
                         examples=[
                             ["What is breast cancer?"],
                             ["What are treatments for cervical cancer?"]
                         ])
    
    return iface.launch(debug = True, share=True)

main()