Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForSeq2SeqLM, T5Tokenizer
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 import time
 import sys
 import traceback
@@ -9,40 +9,32 @@ import traceback
 error_message = ""
 
 # Load the model and tokenizer from Hugging Face
-model_name = "ambrosfitz/history-qa-t5-base-large"
+model_name = "ambrosfitz/history-qa-flan-t5-large"
 try:
     model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-    tokenizer = T5Tokenizer.from_pretrained(model_name, use_fast=False)
+    tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device)
 except Exception as e:
     error_message = f"Error loading model or tokenizer: {str(e)}\n{traceback.format_exc()}"
-    print(error_message)
+print(error_message)
 
 def generate_qa(text, max_length=512):
     try:
-        input_text = f"Generate question: {text}"
+        input_text = f"Generate a history question and answer based on this text: {text}"
         input_ids = tokenizer(input_text, return_tensors="pt", max_length=max_length, truncation=True).input_ids.to(device)
         
         with torch.no_grad():
             outputs = model.generate(input_ids, max_length=max_length, num_return_sequences=1, do_sample=True, temperature=0.7)
-        
-        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
         # Parse the generated text
         parts = generated_text.split("Question: ")
         if len(parts) > 1:
-            qa_parts = parts[1].split("Options:")
+            qa_parts = parts[1].split("Answer: ")
             question = qa_parts[0].strip()
-            
-            options_and_answer = qa_parts[1].split("Correct Answer:")
-            options = options_and_answer[0].strip()
-            
-            answer_and_explanation = options_and_answer[1].split("Explanation:")
-            correct_answer = answer_and_explanation[0].strip()
-            explanation = answer_and_explanation[1].strip() if len(answer_and_explanation) > 1 else "No explanation provided."
-            
-            return f"Question: {question}\n\nOptions: {options}\n\nCorrect Answer: {correct_answer}\n\nExplanation: {explanation}"
+            answer = qa_parts[1].strip() if len(qa_parts) > 1 else "No answer provided."
+            return f"Question: {question}\n\nAnswer: {answer}"
         else:
             return "Unable to generate a proper question and answer. Please try again with a different input."
     except Exception as e:
@@ -63,8 +55,8 @@ try:
         slow_qa,
         chatbot=gr.Chatbot(height=500),
         textbox=gr.Textbox(placeholder="Enter historical text here...", container=False, scale=7),
-        title="History Q&A Generator",
-        description="Enter a piece of historical text, and the model will generate a related question, answer options, correct answer, and explanation.",
+        title="History Q&A Generator (FLAN-T5)",
+        description="Enter a piece of historical text, and the model will generate a related question and answer.",
         theme="soft",
         examples=[
             "The American Revolution was a colonial revolt that took place between 1765 and 1783.",
@@ -79,9 +71,8 @@ try:
     
     if error_message:
         print("Launching interface with error message.")
-        iface.launch(debug=True)
     else:
         print("Launching interface normally.")
-        iface.launch(debug=True)
+    iface.launch(debug=True)
 except Exception as e:
     print(f"An error occurred while creating or launching the interface: {str(e)}\n{traceback.format_exc()}")