Update app.py

app.py

@@ -3,14 +3,6 @@ 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']
@@ -18,26 +10,29 @@ def prepare_context():
     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
+          cleaned_text = text[:objective_index]
+          pubmed_information_cleaned += cleaned_text
         else:
-            pubmed_information_cleaned += text
-    max_length = 1000  # Adjust as needed
+          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},
+      {
+          "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 = f"{pubmed_information_cleaned}\n\n"
     prompt_with_pubmed += tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
 
     # Generate response
@@ -51,7 +46,7 @@ def answer_question(question):
 
 def main():
     """"
-    Initializes a Women Cancer ChatBot interface using Hugging Face models for question answering.
+    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
@@ -63,15 +58,16 @@ def main():
     >>> 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?"]
-                         ])
-    
+                    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)
 
-main()
+if __name__ == "__main__":
+    main()