Update backend.py

backend.py

@@ -1,34 +1,52 @@
-from transformers import T5ForConditionalGeneration, T5Tokenizer
-import torch
+from transformers import T5Tokenizer, T5ForConditionalGeneration
 
-# Load model and tokenizer once
+# Loading the fine-tuned model
 model_name = "./T5base_Question_Generation"
 tokenizer = T5Tokenizer.from_pretrained(model_name)
 model = T5ForConditionalGeneration.from_pretrained(model_name)
 
-def generate_questions(context, question_type, difficulty, num_questions=1):
+def get_question(tag, difficulty, context, answer="", num_questions=1, max_length=150):
     """
-    Generates questions using T5 model based on the provided context, question type, and difficulty.
+    Generate questions using the fine-tuned T5 model
+    
+    Parameters:
+    - tag: Type of question (e.g., "short answer", "multiple choice question", "true or false question")
+    - difficulty: "easy", "medium", "hard"
+    - context: Supporting context or passage
+    - answer: Optional — if you want targeted question generation
+    - num_questions: Number of diverse questions to generate
+    - max_length: Max token length of generated output
+    
+    Returns:
+    - List of generated questions as strings
     """
-    # Build the prompt
-    prompt = f"Generate {difficulty} {question_type} question based on: {context}"
+    # Format input text based on whether answer is provided
+    answer_part = f"[{answer}]" if answer else ""
+    input_text = f"<extra_id_97>{tag} <extra_id_98>{difficulty} <extra_id_99>{answer_part} {context}"
 
     # Tokenize input
-    input_ids = tokenizer.encode(prompt, return_tensors="pt", truncation=True)
+    features = tokenizer([input_text], return_tensors='pt', truncation=True, padding=True)
 
-    # Generate output
+    # Decide generation strategy
     if num_questions == 1:
-        outputs = model.generate(input_ids, max_length=64)
+        output = model.generate(
+            input_ids=features['input_ids'],
+            attention_mask=features['attention_mask'],
+            max_length=max_length,
+            do_sample=False
+        )
     else:
-        outputs = model.generate(
-            input_ids,
-            max_length=64,
+        output = model.generate(
+            input_ids=features['input_ids'],
+            attention_mask=features['attention_mask'],
+            max_length=max_length,
             do_sample=True,
-            top_k=50,
             top_p=0.95,
-            num_return_sequences=num_questions,
+            top_k=50,
+            num_return_sequences=num_questions
         )
 
-    # Decode and return questions
-    questions = [tokenizer.decode(output, skip_special_tokens=True) for output in outputs]
+    # Decode questions
+    questions = [tokenizer.decode(out, skip_special_tokens=True) for out in output]
     return questions
+