t5-history-MC / app2.py
ambrosfitz's picture
Rename app.py to app2.py
ba2ecb5 verified
import gradio as gr
import torch
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
import time
import sys
import traceback
# Global variables to store error information
error_message = ""
# Load the model and tokenizer from Hugging Face
model_name = "ambrosfitz/history-qa-flan-t5-large"
try:
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
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)
def generate_qa(text, max_length=512):
try:
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)
# Parse the generated text
parts = generated_text.split("Question: ")
if len(parts) > 1:
qa_parts = parts[1].split("Answer: ")
question = qa_parts[0].strip()
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:
return f"An error occurred: {str(e)}\n{traceback.format_exc()}"
def slow_qa(message, history):
try:
full_response = generate_qa(message)
for i in range(len(full_response)):
time.sleep(0.01)
yield full_response[:i+1]
except Exception as e:
yield f"An error occurred: {str(e)}\n{traceback.format_exc()}"
# Create and launch the Gradio interface
try:
iface = gr.ChatInterface(
slow_qa,
chatbot=gr.Chatbot(height=500),
textbox=gr.Textbox(placeholder="Enter historical text here...", container=False, scale=7),
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.",
"World War II was a global conflict that lasted from 1939 to 1945, involving many of the world's nations.",
"The Renaissance was a period of cultural, artistic, political, and economic revival following the Middle Ages."
],
cache_examples=False,
retry_btn="Regenerate",
undo_btn="Remove last",
clear_btn="Clear",
)
if error_message:
print("Launching interface with error message.")
else:
print("Launching interface normally.")
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()}")