Update app.py

app.py

@@ -1,32 +1,21 @@
+import spaces
 import torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 MODEL_NAME = "ma4389/LFM2-DPO"
 
-print("Loading model...")
-
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    device_map="auto" if torch.cuda.is_available() else None,
-)
-
-model.eval()
-
 SYSTEM_PROMPT = """
 You are an expert educational question generation assistant.
 
-Generate high-quality questions ONLY from the provided paragraph.
+Generate questions ONLY from the supplied paragraph.
 
 Rules:
-- Use only the given paragraph.
-- Do not hallucinate.
-- Do not use outside knowledge.
-- Cover different concepts.
-- Return only the requested format.
+- Use ONLY the provided paragraph.
+- Never hallucinate.
+- Never repeat questions.
+- Each question must test a different concept.
+- Return ONLY the requested format.
 """
 
 PROMPTS = {
@@ -36,10 +25,18 @@ Paragraph:
 
 Generate exactly {num} multiple-choice questions.
 
-Requirements:
-- Four choices (A, B, C, D)
-- One correct answer
-- No explanations
+Rules:
+- Use ONLY the paragraph.
+- Every question must test a different concept.
+- Never repeat questions.
+- Each question must have exactly four choices:
+A)
+B)
+C)
+D)
+- The correct answer should vary naturally between A, B, C and D.
+- Wrong choices should be realistic.
+- Do not explain the answer.
 
 Format:
 
@@ -50,7 +47,7 @@ B) ...
 C) ...
 D) ...
 
-Answer: A
+Answer: C
 """,
 
     "True / False": """
@@ -59,6 +56,12 @@ Paragraph:
 
 Generate exactly {num} True/False questions.
 
+Rules:
+- Use ONLY the paragraph.
+- Never repeat questions.
+- Mix True and False naturally.
+- Do not explain.
+
 Format:
 
 1. Statement
@@ -66,26 +69,55 @@ Format:
 Answer: True
 """,
 
-    "WH Questions": """
+    "Essay": """
 Paragraph:
 {context}
 
-Generate exactly {num} short-answer questions.
+Generate exactly {num} essay questions.
 
-Requirements:
-- Begin with Who, What, When, Where, Why, or How whenever appropriate.
-- Each question must have its answer.
+Rules:
+- Use ONLY the paragraph.
+- Every question should test a different concept.
+- Never repeat questions.
+- Each answer should contain 3-8 complete sentences.
+- Do not invent information.
 
 Format:
 
 1. Question
 
-Answer: ...
+Answer:
+Detailed answer...
 """
 }
 
+# Lazy loading
+model = None
+tokenizer = None
+
+
+def load_model():
+    global model, tokenizer
+
+    if model is None:
+        print("Loading model...")
 
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
+            dtype=torch.float16,
+            device_map="auto",
+        )
+
+        model.eval()
+
+
+@spaces.GPU
 def ask(prompt):
+
+    load_model()
+
     messages = [
         {"role": "system", "content": SYSTEM_PROMPT},
         {"role": "user", "content": prompt},
@@ -95,21 +127,23 @@ def ask(prompt):
         messages,
         tokenize=True,
         add_generation_prompt=True,
-        return_tensors="pt",
         return_dict=True,
+        return_tensors="pt",
     )
 
     inputs = {k: v.to(model.device) for k, v in inputs.items()}
 
     with torch.no_grad():
+
         outputs = model.generate(
             **inputs,
-            max_new_tokens=512,
-            temperature=0.3,
+            max_new_tokens=700,
+            temperature=0.4,
             top_p=0.9,
+            top_k=50,
             do_sample=True,
-            repetition_penalty=1.15,
-            no_repeat_ngram_size=3,
+            repetition_penalty=1.3,
+            no_repeat_ngram_size=5,
             pad_token_id=tokenizer.eos_token_id,
             eos_token_id=tokenizer.eos_token_id,
         )
@@ -120,45 +154,46 @@ def ask(prompt):
     )
 
 
-def generate_questions(paragraph, question_type, num_questions):
+def generate(paragraph, qtype, num):
+
     if not paragraph.strip():
         return "Please enter a paragraph."
 
-    prompt = PROMPTS[question_type].format(
+    prompt = PROMPTS[qtype].format(
         context=paragraph,
-        num=num_questions,
+        num=num,
     )
 
     return ask(prompt)
 
 
 demo = gr.Interface(
-    fn=generate_questions,
+    fn=generate,
     inputs=[
         gr.Textbox(
-            lines=12,
             label="Paragraph",
+            lines=15,
             placeholder="Paste any paragraph here..."
         ),
         gr.Radio(
-            ["MCQ", "True / False", "WH Questions"],
+            ["MCQ", "True / False", "Essay"],
             value="MCQ",
-            label="Question Type",
+            label="Question Type"
         ),
         gr.Slider(
             minimum=1,
             maximum=10,
             value=5,
             step=1,
-            label="Number of Questions",
+            label="Number of Questions"
         ),
     ],
     outputs=gr.Textbox(
-        lines=20,
         label="Generated Questions",
+        lines=25,
     ),
     title="📚 AI Question Generator",
-    description="Generate MCQ, True/False, or WH questions from any paragraph using a fine-tuned LFM2-DPO model.",
+    description="Generate MCQ, True/False and Essay questions from any paragraph using a fine-tuned LFM2-DPO model.",
 )
 
 if __name__ == "__main__":