Update backend.py

backend.py

@@ -5,7 +5,7 @@ model_name = "./T5base_Question_Generation"
 tokenizer = T5Tokenizer.from_pretrained(model_name)
 model = T5ForConditionalGeneration.from_pretrained(model_name)
 
-def get_question(tag, difficulty, context, answer="", num_questions=1, max_length=150):
+def get_question(tag, difficulty, context, answer="", num_questions=1, use_beam_search=False, num_beams=3, max_length=150):
     """
     Generate questions using the fine-tuned T5 model
     
@@ -29,12 +29,21 @@ def get_question(tag, difficulty, context, answer="", num_questions=1, max_lengt
 
     # Decide generation strategy
     if num_questions == 1:
-        output = model.generate(
-            input_ids=features['input_ids'],
-            attention_mask=features['attention_mask'],
-            max_length=max_length,
-            do_sample=False
-        )
+        if use_beam_search:
+            output = model.generate(
+                input_ids=features['input_ids'],
+                attention_mask=features['attention_mask'],
+                max_length=max_length,
+                num_beams=num_beams,
+                early_stopping=False
+            )
+        else:
+            output = model.generate(
+                input_ids=features['input_ids'],
+                attention_mask=features['attention_mask'],
+                max_length=max_length,
+                do_sample=False
+            )
     else:
         output = model.generate(
             input_ids=features['input_ids'],