Update app.py

app.py

@@ -1,16 +1,21 @@
 # app.py
 import gradio as gr
 from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
 
 # ====== Load AI models ======
 # Summarization model (BART)
 summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
 
-# Question generation model (T5 small, CPU-friendly)
-# Manually load tokenizer to avoid SentencePiece fast tokenizer issues
+# ====== Question / Flashcard model (T5 small) ======
 tokenizer = AutoTokenizer.from_pretrained("valhalla/t5-small-qg-hl", use_fast=False)
 model = AutoModelForSeq2SeqLM.from_pretrained("valhalla/t5-small-qg-hl")
-question_generator = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
+
+# ====== Helper function for manual generation ======
+def generate_questions(text, max_length=128):
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+    outputs = model.generate(**inputs, max_length=max_length)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 # ====== Summarizer Function ======
 def summarize_text(text, mode):
@@ -31,15 +36,13 @@ def summarize_text(text, mode):
 def generate_quiz(text):
     if not text.strip():
         return "⚠️ Please provide some notes!"
-    output = question_generator(f"generate questions: {text}", max_length=256)
-    return output[0]['generated_text']
+    return generate_questions(f"generate questions: {text}")
 
 # ====== Flashcards Generator ======
 def generate_flashcards(text):
     if not text.strip():
         return "⚠️ Please provide some notes!"
-    output = question_generator(f"generate flashcards: {text}", max_length=256)
-    return output[0]['generated_text']
+    return generate_questions(f"generate flashcards: {text}")
 
 # ====== Gradio Interface ======
 with gr.Blocks() as demo: