Update src/generation.py

src/generation.py

@@ -15,6 +15,7 @@ CONTRACT -- do not change:
     answer(question: str, chunks: list[Chunk], best_score: float) -> Answer
 =============================================================================
 """
+from transformers import pipeline
 
 from config import FALLBACK_MESSAGE, GENERATION_MODEL, MAX_NEW_TOKENS, RELEVANCE_THRESHOLD
 from src.types import Answer, Chunk
@@ -31,6 +32,53 @@ Answer:"""
 # instruction-tuned but small, so phrasing matters a lot. Record what you
 # tried and what changed; prompt iteration is your results-paper material.
 
+# flan-t5-base's encoder handles ~512 tokens total. Leave headroom for the
+# template text + question, so budget the context block conservatively.
+# If this keeps truncating chunks in practice, that's a signal to ask
+# Suyash for fewer/shorter chunks rather than raising this further.
+MAX_CONTEXT_TOKENS = 400
+
+_generator = None  # module-level cache -- loaded once, reused across calls
+
+
+def _get_generator():
+    """Lazily load and cache the generation pipeline.
+
+    Reloading the model per question is unusably slow (multi-second load
+    every call), so this is only ever done once per process.
+    """
+    global _generator
+    if _generator is None:
+        _generator = pipeline("text2text-generation", model=GENERATION_MODEL)
+    return _generator
+
+
+def _build_context(chunks: list[Chunk], tokenizer) -> str:
+    """Join chunk texts into the context block, truncating to fit the
+    model's context window if necessary.
+
+    Chunks are added in the order given (i.e. Suyash's ranking) and we stop
+    adding once we'd exceed MAX_CONTEXT_TOKENS, rather than truncating mid
+    chunk -- a partial chunk is more likely to mislead the model than a
+    dropped low-ranked one.
+    """
+    parts: list[str] = []
+    used_tokens = 0
+
+    for chunk in chunks:
+        chunk_tokens = len(tokenizer.encode(chunk.text))
+        if used_tokens + chunk_tokens > MAX_CONTEXT_TOKENS:
+            if not parts:
+                # Even the single best chunk is too long -- truncate it
+                # directly rather than returning empty context.
+                truncated_ids = tokenizer.encode(chunk.text)[:MAX_CONTEXT_TOKENS]
+                parts.append(tokenizer.decode(truncated_ids, skip_special_tokens=True))
+            break
+        parts.append(chunk.text)
+        used_tokens += chunk_tokens
+
+    return "\n\n".join(parts)
+
 
 def answer(question: str, chunks: list[Chunk], best_score: float = 1.0) -> Answer:
     """Generate an answer grounded in the retrieved chunks.
@@ -50,29 +98,13 @@ def answer(question: str, chunks: list[Chunk], best_score: float = 1.0) -> Answe
     if not chunks or best_score < RELEVANCE_THRESHOLD:
         return Answer(text=FALLBACK_MESSAGE, chunks=[], in_scope=False)
 
-    # =====================================================================
-    # AKSHAY: YOUR CODE GOES HERE
-    #
-    # Suggested steps:
-    #   1. Lazily load GENERATION_MODEL via transformers pipeline
-    #      ("text2text-generation" for flan-t5). Cache it at module level
-    #      -- reloading per question is unusably slow.
-    #   2. Join the chunk texts into the {context} slot of PROMPT_TEMPLATE.
-    #      Mind the context window: flan-t5-base handles ~512 tokens, so
-    #      TOP_K chunks of CHUNK_SIZE may need truncating. Coordinate with
-    #      Suyash if you need him to return fewer or shorter chunks.
-    #   3. Generate with MAX_NEW_TOKENS.
-    #   4. Return Answer(text=..., chunks=chunks, in_scope=True). Always
-    #      pass the chunks back -- that is what makes citation possible.
-    #
-    # Watch for: the model answering from its own pretraining rather than
-    # the context. Test with a deliberately wrong context and confirm it
-    # follows the context, not its own memory.
-    # =====================================================================
-    _ = (GENERATION_MODEL, MAX_NEW_TOKENS, PROMPT_TEMPLATE)  # remove when implemented
-
-    return Answer(
-        text=f"[STUB ANSWER] {chunks[0].text}",
-        chunks=chunks,
-        in_scope=True,
-    )
+    generator = _get_generator()
+    tokenizer = generator.tokenizer
+
+    context = _build_context(chunks, tokenizer)
+    prompt = PROMPT_TEMPLATE.format(context=context, question=question)
+
+    result = generator(prompt, max_new_tokens=MAX_NEW_TOKENS, do_sample=False)
+    text = result[0]["generated_text"].strip()
+
+    return Answer(text=text, chunks=chunks, in_scope=True)