Update app.py

app.py

@@ -1,33 +1,34 @@
 import gradio as gr
-from transformers import T5Tokenizer, T5ForConditionalGeneration
-import torch
+from pprint import pprint
+from lmqg import TransformersQG
 
-model_name = "t5-small"
+# Initialize model (CPU-only)
+model = TransformersQG(model='lmqg/t5-base-squad-qg', model_ae='lmqg/t5-base-squad-ae')
 
-tokenizer = T5Tokenizer.from_pretrained(model_name)
-model = T5ForConditionalGeneration.from_pretrained(model_name)
-model.eval()  # set model to evaluation mode
+# Function to chunk text into smaller parts
+def chunk_text(text, chunk_size=450):
+    tokenizer = model.tokenizer
+    tokens = tokenizer.encode(text)
+    chunks = [tokens[i:i + chunk_size] for i in range(0, len(tokens), chunk_size)]
+    return chunks
 
-device = torch.device("cpu")  # explicitly set device to CPU
-model.to(device)
-
-def generate_text(input_text):
-    input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
-    outputs = model.generate(
-        input_ids,
-        max_length=100,
-        num_beams=5,
-        early_stopping=True
-    )
-    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return result
+# Function to process each chunk and generate QA pairs
+def generate_qa_for_chunks(text):
+    chunks = chunk_text(text)
+    qa_pairs = []
+    for chunk in chunks:
+        chunk_text = model.tokenizer.decode(chunk, skip_special_tokens=True)
+        qa = model.generate_qa(chunk_text)
+        qa_pairs.extend(qa)
+    return "\n\n".join([f"Q: {item['question']}\nA: {item['answer']}" for item in qa_pairs])
 
+# Gradio UI
 demo = gr.Interface(
-    fn=generate_text,
-    inputs=gr.Textbox(lines=5, label="Input Text"),
-    outputs=gr.Textbox(lines=5, label="Generated Text"),
-    title="Simple T5 Text Generator (CPU)",
-    description="Enter some text, T5 will generate a continuation or answer."
+    fn=generate_qa_for_chunks,
+    inputs=gr.Textbox(lines=20, label="Input Text (SAP Note or Paragraph)"),
+    outputs=gr.Textbox(lines=30, label="Generated QA Pairs"),
+    title="Question Generator (LMQG - T5)",
+    description="Paste text to generate question-answer pairs using lmqg/t5-base-squad-qg"
 )
 
 if __name__ == "__main__":