Update app.py

app.py

@@ -6,7 +6,6 @@ from langchain.text_splitter import CharacterTextSplitter
 from langchain_community.vectorstores import FAISS
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_core.prompts import PromptTemplate
-from langchain_community.llms import HuggingFacePipeline
 
 # Hugging Face Transformers
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
@@ -15,7 +14,7 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 # ---------------- Load LLM with fallback ----------------
 def load_llm():
     model_ids = [
-        "google/flan-t5-base",  # Start with base model (more likely to work)
+        "google/flan-t5-base",
         "google/flan-t5-small",
         "google/flan-t5-large"
     ]
@@ -26,7 +25,7 @@ def load_llm():
             tokenizer = AutoTokenizer.from_pretrained(model_id)
             model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
             
-            # Create pipeline without return_full_text parameter
+            # Create pipeline directly without LangChain wrapper
             pipe = pipeline(
                 "text2text-generation",
                 model=model,
@@ -34,17 +33,15 @@ def load_llm():
                 max_length=512
             )
             print(f"✅ Successfully loaded model: {model_id}")
-            
-            # Create HuggingFacePipeline without any extra parameters
-            return HuggingFacePipeline(pipeline=pipe)
+            return pipe
         except Exception as e:
             print(f"⚠️ Failed to load {model_id}: {e}")
             continue
 
-    raise RuntimeError("❌ No model could be loaded. Please check Hugging Face space resources.")
+    raise RuntimeError("❌ No model could be loaded.")
 
 
-llm = load_llm()
+llm_pipeline = load_llm()
 
 
 # ---------------- Process PDF ----------------
@@ -88,13 +85,17 @@ def ask_question(pdf_files, question):
         context = "\n".join([doc.page_content for doc in docs])
 
         # Prompt template
-        prompt = PromptTemplate(
-            input_variables=["context", "question"],
-            template="Answer the question based on the following context:\n\nContext: {context}\n\nQuestion: {question}\nAnswer:"
-        )
+        prompt_template = f"""Answer the question based on the following context:
+
+Context: {context}
 
-        final_prompt = prompt.format(context=context, question=question)
-        response = llm.invoke(final_prompt)
+Question: {question}
+
+Answer:"""
+
+        # Use the pipeline directly
+        result = llm_pipeline(prompt_template, max_length=512, do_sample=False)
+        response = result[0]['generated_text'].strip()
 
         return response if response else "⚠️ No answer generated. Try another question."
 
@@ -126,13 +127,4 @@ with gr.Blocks() as demo:
     submit_btn = gr.Button("Ask Question", variant="primary")
     submit_btn.click(fn=ask_question, inputs=[pdf_input, question_input], outputs=output)
 
-    # Add examples
-    gr.Examples(
-        examples=[
-            ["What is the main topic of this document?", "Summarize the key points"],
-            ["What are the main findings?", "Who are the authors?"]
-        ],
-        inputs=question_input
-    )
-
 demo.launch()