Update app.py

app.py

@@ -70,40 +70,78 @@ def generate(passage: str, language_code: str, n: int):
 
     model, tok = _load()
     lang_name = LANG_NAMES.get(language_code, language_code)
+
     slots = ", ".join(['"..."'] * int(n))
-    user = TEACHER_TEMPLATE.format(n=int(n), language=lang_name, slots=slots, passage=passage)
+    user = TEACHER_TEMPLATE.format(
+        n=int(n),
+        language=lang_name,
+        slots=slots,
+        passage=passage,
+    )
+
     messages = [
         {"role": "system", "content": TEACHER_SYSTEM},
         {"role": "user", "content": user},
     ]
-    input_ids = tok.apply_chat_template(
-        messages, tokenize=True, add_generation_prompt=True, return_tensors="pt"
-    ).to(model.device)
+
+    # Build model inputs
+    inputs = tok.apply_chat_template(
+        messages,
+        tokenize=True,
+        add_generation_prompt=True,
+        return_tensors="pt",
+        return_dict=True,
+    )
+
+    inputs = {k: v.to(model.device) for k, v in inputs.items()}
 
     with torch.no_grad():
         out = model.generate(
-            input_ids, max_new_tokens=250, do_sample=False,
-            pad_token_id=tok.pad_token_id or tok.eos_token_id,
+            **inputs,
+            max_new_tokens=250,
+            do_sample=False,
+            pad_token_id=tok.pad_token_id
+            if tok.pad_token_id is not None
+            else tok.eos_token_id,
         )
-    text = tok.decode(out[0][input_ids.shape[1]:], skip_special_tokens=True)
+
+    prompt_len = inputs["input_ids"].shape[1]
+
+    text = tok.decode(
+        out[0][prompt_len:],
+        skip_special_tokens=True,
+    )
 
     try:
-        obj = json.loads(text[text.index("{"): text.rindex("}") + 1])
-        questions = [q for q in obj.get("questions", []) if isinstance(q, str) and q.strip()]
+        start = text.index("{")
+        end = text.rindex("}") + 1
+        obj = json.loads(text[start:end])
+        questions = [
+            q
+            for q in obj.get("questions", [])
+            if isinstance(q, str) and q.strip()
+        ]
     except (ValueError, json.JSONDecodeError):
         questions = []
 
     note = ""
     if language_code in WEAK_LANGS:
-        note = (f"\n\n⚠️ {lang_name} is one of this model's lower-confidence languages "
-                f"(see the [model card](https://huggingface.co/{MODEL_ID}) for benchmark numbers).")
+        note = (
+            f"\n\n⚠️ {lang_name} is one of this model's lower-confidence languages "
+            f"(see the [model card](https://huggingface.co/{MODEL_ID}) for benchmark numbers)."
+        )
 
     if not questions:
-        return text, f"⚠️ Couldn't parse valid questions from the model's output. Raw output shown alongside.{note}"
+        return (
+            text,
+            f"⚠️ Couldn't parse valid questions from the model's output. Raw output shown alongside.{note}",
+        )
 
-    formatted = "\n".join(f"{i+1}. {q}" for i, q in enumerate(questions))
-    return formatted, ("✅ Generated." + note) if note else "✅ Generated."
+    formatted = "\n".join(
+        f"{i + 1}. {q}" for i, q in enumerate(questions)
+    )
 
+    return formatted, ("✅ Generated." + note) if note else "✅ Generated."
 
 with gr.Blocks(title="mist-qg-1.5b — multilingual question generator") as demo:
     gr.Markdown(