Update app.py

app.py

@@ -67,20 +67,27 @@ class LangChainAgentWrapper:
     def __init__(self):
         print("Initializing LangChainAgentWrapper...")
         
-        # Switched to a smaller, CPU-friendly instruction-tuned model
         model_id = "google/flan-t5-base"
         
         try:
-            hf_auth_token = os.getenv("HF_TOKEN")
             print(f"Loading model pipeline for: {model_id}")
             
-            # For FLAN-T5, we use the "text2text-generation" task.
-            # We also remove quantization as it's not needed for this smaller model.
-            llm_pipeline = transformers.pipeline(
-                "text2text-generation", # <<< IMPORTANT: Changed task for T5 models
-                model=model_id,
-                device_map="auto"
+            # --- MODIFICATION: Use the custom pipeline class ---
+            # Load the tokenizer first
+            tokenizer = transformers.AutoTokenizer.from_pretrained(model_id)
+            # Load the model
+            model = transformers.AutoModelForSeq2SeqLM.from_pretrained(model_id)
+
+            # Create an instance of our custom pipeline
+            llm_pipeline = FlanT5Pipeline(
+                task="text2text-generation",
+                model=model,
+                tokenizer=tokenizer,
+                device_map="auto",
+                max_new_tokens=512 # Add max_new_tokens to control output length
             )
+            # --- END MODIFICATION ---
+
             print("Model pipeline loaded successfully.")
             
             # Wrap the pipeline in a LangChain LLM object