Update generator_engine.py

generator_engine.py

@@ -1,82 +1,129 @@
 from __future__ import annotations
 
-
-def normalize_category(category: str | None) -> str:
-    c = (category or "").strip().lower()
-
-    if c in {"quantitative", "quant", "q", "math"}:
-        return "Quantitative"
-    if c in {"datainsight", "data_insight", "data insight", "di", "data"}:
-        return "DataInsight"
-    if c in {"verbal", "v"}:
-        return "Verbal"
-    if c in {"general", "", "unknown", "none", "null"}:
-        return "General"
-
-    return category or "General"
-
-
-def classify_question(question_text: str, category: str | None = None) -> dict:
-    q = (question_text or "").lower()
-    normalized = normalize_category(category)
-
-    if normalized == "Quantitative":
-        if ("percent" in q or "%" in q) and any(
-            k in q for k in ["then", "after", "followed by", "successive", "increase", "decrease", "discount"]
-        ):
-            return {"category": normalized, "topic": "percent", "type": "successive_percent"}
-
-        if "percent" in q or "%" in q:
-            return {"category": normalized, "topic": "percent", "type": "percent_change"}
-
-        if "ratio" in q or ":" in q:
-            return {"category": normalized, "topic": "ratio", "type": "ratio_total"}
-
-        if "probability" in q or "chosen at random" in q:
-            return {"category": normalized, "topic": "probability", "type": "simple_probability"}
-
-        if "divisible" in q or "remainder" in q or "mod" in q:
-            return {"category": normalized, "topic": "number_theory", "type": "remainder_or_divisibility"}
-
-        if "|" in q:
-            return {"category": normalized, "topic": "algebra", "type": "absolute_value"}
-
-        if any(k in q for k in ["circle", "radius", "circumference", "triangle", "perimeter", "area"]):
-            return {"category": normalized, "topic": "geometry", "type": "geometry"}
-
-        if any(k in q for k in ["average", "mean", "median"]):
-            return {"category": normalized, "topic": "statistics", "type": "average"}
-
-        if "sequence" in q:
-            return {"category": normalized, "topic": "sequence", "type": "sequence"}
-
-        if "=" in q:
-            return {"category": normalized, "topic": "algebra", "type": "equation"}
-
-        return {"category": normalized, "topic": "quant", "type": "general"}
-
-    if normalized == "DataInsight":
-        if "percent" in q or "%" in q:
-            return {"category": normalized, "topic": "percent", "type": "percent_change"}
-        if any(k in q for k in ["mean", "median", "distribution"]):
-            return {"category": normalized, "topic": "statistics", "type": "distribution"}
-        if any(k in q for k in ["correlation", "scatter", "trend", "table", "chart"]):
-            return {"category": normalized, "topic": "data", "type": "correlation_or_graph"}
-        return {"category": normalized, "topic": "data", "type": "general"}
-
-    if normalized == "Verbal":
-        if "meaning" in q or "definition" in q:
-            return {"category": normalized, "topic": "vocabulary", "type": "definition"}
-        if "grammatically" in q or "sentence correction" in q:
-            return {"category": normalized, "topic": "grammar", "type": "sentence_correction"}
-        if "argument" in q or "author" in q:
-            return {"category": normalized, "topic": "reasoning", "type": "argument_analysis"}
-        return {"category": normalized, "topic": "verbal", "type": "general"}
-
-    if any(k in q for k in ["percent", "%", "ratio", "remainder", "divisible", "probability", "circle", "triangle", "="]):
-        return classify_question(question_text, "Quantitative")
-
-    if any(k in q for k in ["table", "chart", "scatter", "trend", "distribution"]):
-        return classify_question(question_text, "DataInsight")
-
-    return {"category": "General", "topic": "unknown", "type": "unknown"}
+from typing import List, Optional
+
+try:
+    from transformers import pipeline
+except Exception:
+    pipeline = None
+
+from models import RetrievedChunk
+
+
+class GeneratorEngine:
+    def __init__(self, model_name: str = "google/flan-t5-small"):
+        self.model_name = model_name
+        self.pipe = None
+
+        if pipeline is not None:
+            try:
+                self.pipe = pipeline("text2text-generation", model=model_name)
+            except Exception:
+                self.pipe = None
+
+    def available(self) -> bool:
+        return self.pipe is not None
+
+    def _notes_block(self, retrieval_context: List[RetrievedChunk]) -> str:
+        if not retrieval_context:
+            return ""
+        lines = []
+        for chunk in retrieval_context[:3]:
+            text = (chunk.text or "").strip().replace("\n", " ")
+            if len(text) > 220:
+                text = text[:217].rstrip() + "…"
+            lines.append(f"- {chunk.topic}: {text}")
+        return "\n".join(lines)
+
+    def _template_fallback(
+        self,
+        user_text: str,
+        question_text: Optional[str],
+        topic: str,
+        intent: str,
+        retrieval_context: Optional[List[RetrievedChunk]] = None,
+    ) -> str:
+        notes = self._notes_block(retrieval_context or [])
+
+        if intent == "hint":
+            base = "Start by identifying the exact relationship between the quantities before doing any arithmetic."
+        elif intent in {"instruction", "method"}:
+            base = "Translate the wording into an equation, ratio, or percent relationship, then solve one step at a time."
+        elif intent in {"walkthrough", "step_by_step", "explain", "concept"}:
+            base = "First identify what the question is asking, then map the values into the correct quantitative structure, and only then compute."
+        else:
+            base = "This does not match a strong solver rule yet, so begin by identifying the target quantity and the relationship connecting the numbers."
+
+        if notes:
+            return f"{base}\n\nRelevant notes:\n{notes}"
+        return base
+
+    def _build_prompt(
+        self,
+        user_text: str,
+        question_text: Optional[str],
+        topic: str,
+        intent: str,
+        retrieval_context: Optional[List[RetrievedChunk]] = None,
+    ) -> str:
+        question = (question_text or user_text or "").strip()
+        notes = self._notes_block(retrieval_context or [])
+
+        prompt = [
+            "You are a concise GMAT tutor.",
+            f"Topic: {topic or 'general'}",
+            f"Intent: {intent or 'answer'}",
+            "",
+            f"Question: {question}",
+        ]
+
+        if notes:
+            prompt.extend(["", "Relevant teaching notes:", notes])
+
+        prompt.extend(
+            [
+                "",
+                "Respond briefly and clearly.",
+                "If the problem is not fully solvable from the parse, give the next best method step.",
+                "Do not invent facts.",
+            ]
+        )
+
+        return "\n".join(prompt)
+
+    def generate(
+        self,
+        user_text: str,
+        question_text: Optional[str] = None,
+        topic: str = "",
+        intent: str = "answer",
+        retrieval_context: Optional[List[RetrievedChunk]] = None,
+        chat_history=None,
+        max_new_tokens: int = 96,
+        **kwargs,
+    ) -> Optional[str]:
+        prompt = self._build_prompt(
+            user_text=user_text,
+            question_text=question_text,
+            topic=topic,
+            intent=intent,
+            retrieval_context=retrieval_context or [],
+        )
+
+        if self.pipe is not None:
+            try:
+                out = self.pipe(prompt, max_new_tokens=max_new_tokens, do_sample=False)
+                if out and isinstance(out, list):
+                    text = str(out[0].get("generated_text", "")).strip()
+                    if text:
+                        return text
+            except Exception:
+                pass
+
+        return self._template_fallback(
+            user_text=user_text,
+            question_text=question_text,
+            topic=topic,
+            intent=intent,
+            retrieval_context=retrieval_context or [],
+        )