| from __future__ import annotations |
|
|
| 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 [], |
| ) |
