Update mini_ladder.py

mini_ladder.py

@@ -1,12 +1,22 @@
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer
 
-# A second pipeline for self-critique (using a lighter model for demonstration)
-CRITIQUE_MODEL = "gpt2"  # This can be replaced with another model as needed
-critique_pipeline = pipeline("text-generation", model=CRITIQUE_MODEL)
+# ------------------------------
+# 1) CRITIQUE MODEL & TOKENIZER
+# ------------------------------
+# Using GPT-2 for self-critique
+CRITIQUE_MODEL_NAME = "gpt2"
+critique_pipeline = pipeline("text-generation", model=CRITIQUE_MODEL_NAME)
+critique_tokenizer = AutoTokenizer.from_pretrained(CRITIQUE_MODEL_NAME)
 
+# GPT-2 typically has a max context length of 1024 tokens
+GPT2_MAX_CONTEXT = 1024
+
+# ------------------------------
+# 2) SUB-QUESTION GENERATION
+# ------------------------------
 def generate_sub_questions(main_query: str):
     """
-    Naively generates sub-questions for the given main query.
+    Naive approach to generating sub-questions.
     """
     return [
         f"1) What are common causes of {main_query}?",
@@ -14,11 +24,15 @@ def generate_sub_questions(main_query: str):
         f"3) What are non-pharmacological approaches to {main_query}?"
     ]
 
+# ------------------------------
+# 3) SELF-CRITIQUE & REFINEMENT
+# ------------------------------
 def self_critique_and_refine(query: str, initial_answer: str, docs: list):
     """
-    Critiques the initial answer and refines it if necessary.
+    1) Critique the initial answer (GPT-2).
+    2) If needed, refine using the original BioGPT pipeline.
     """
-    # Step 1: Generate a critique using a critique prompt
+    # A) Construct the critique prompt
     critique_prompt = (
         f"The following is an answer to the question '{query}'. "
         "Evaluate its correctness, clarity, and completeness. "
@@ -26,14 +40,24 @@ def self_critique_and_refine(query: str, initial_answer: str, docs: list):
         f"ANSWER:\n{initial_answer}\n\n"
         "CRITIQUE:"
     )
-    critique_gen = critique_pipeline(critique_prompt, max_new_tokens=80, truncation=True)
+
+    # B) Truncate the critique prompt to fit GPT-2’s max context
+    truncated_critique_prompt = _truncate_prompt_for_gpt2(critique_prompt)
+
+    # C) Generate the critique
+    critique_gen = critique_pipeline(
+        truncated_critique_prompt,
+        max_new_tokens=80,   # how many tokens to generate for the critique
+        truncation=True      # ensure we don't exceed the final length
+    )
     if critique_gen and isinstance(critique_gen, list):
         critique_text = critique_gen[0]["generated_text"]
     else:
         critique_text = "No critique generated."
 
-    # Step 2: If the critique suggests issues, refine the answer using the original QA pipeline.
+    # D) If critique suggests issues, refine using BioGPT
     if any(word in critique_text.lower() for word in ["missing", "incomplete", "incorrect", "lacks"]):
+        # Build a refine prompt that includes docs
         refine_prompt = (
             f"Question: {query}\n"
             f"Current Answer: {initial_answer}\n"
@@ -43,8 +67,11 @@ def self_critique_and_refine(query: str, initial_answer: str, docs: list):
             + "\n\n".join(docs)
             + "\nREFINED ANSWER:"
         )
-        # Import the qa_pipeline from backend to reuse it (local import to avoid circular dependencies)
-        from backend import qa_pipeline
+
+        # If BioGPT has similar context limits, you can truncate here too
+        # e.g., refine_prompt = _truncate_prompt_for_biogpt(refine_prompt)
+
+        from backend import qa_pipeline  # Import to avoid circular references
         refined_gen = qa_pipeline(refine_prompt, max_new_tokens=120, truncation=True)
         if refined_gen and isinstance(refined_gen, list):
             refined_answer = refined_gen[0]["generated_text"]
@@ -54,3 +81,17 @@ def self_critique_and_refine(query: str, initial_answer: str, docs: list):
         refined_answer = initial_answer
 
     return refined_answer, critique_text
+
+# ------------------------------
+# 4) HELPER: GPT-2 TRUNCATION
+# ------------------------------
+def _truncate_prompt_for_gpt2(prompt_text: str) -> str:
+    """
+    Token-level truncation to ensure the prompt doesn't exceed GPT-2’s 1024-token limit.
+    """
+    tokens = critique_tokenizer.encode(prompt_text, add_special_tokens=False)
+    if len(tokens) > GPT2_MAX_CONTEXT:
+        # Keep the first 1024 tokens
+        tokens = tokens[:GPT2_MAX_CONTEXT]
+    truncated_text = critique_tokenizer.decode(tokens, skip_special_tokens=True)
+    return truncated_text