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

# Function to prepare context
def prepare_context():
    df = pd.read_csv("splitted_df_jo.csv")
    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
    return pubmed_information_cleaned[:max_length]

# Function to generate answer
def answer_question(question):
    pubmed_information_cleaned = prepare_context()
    model_name = "jrocha/tiny_llama"
    model = AutoModelForCausalLM.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    # 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"
    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 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="Cancer ChatBot",
                    description="How can I help you?",
                    examples=[
                        ["What is prostate cancer?"],
                        ["What are treatments for cervical cancer?"]
                    ])

    return iface.launch(debug = True, share=True)

if __name__ == "__main__":
    main()