Update app.py

app.py

@@ -12,15 +12,15 @@ from huggingface_hub import InferenceClient
 
 HF_TOKEN = os.getenv("HF_TOKEN", None)
 
-# Free IBM Granite model on Hugging Face
-GRANITE_MODEL = "ibm-granite/granite-3.2-8b-instruct"
+# Free, stable, HF-compatible model
+LLM_MODEL = "mistralai/Mistral-7B-Instruct-v0.2"
 
-# Local embedding model (fast + CPU friendly)
+# Local embedding model
 EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
 
-# Load models once per process
+# Load models
 embedding_model = SentenceTransformer(EMBEDDING_MODEL)
-llm = InferenceClient(model=GRANITE_MODEL, token=HF_TOKEN)
+llm = InferenceClient(model=LLM_MODEL, token=HF_TOKEN)
 
 
 # ============================================================
@@ -28,7 +28,6 @@ llm = InferenceClient(model=GRANITE_MODEL, token=HF_TOKEN)
 # ============================================================
 
 def read_pdf(file_obj):
-    """Extract raw text from a PDF file."""
     reader = PdfReader(file_obj)
     pages = []
     for page in reader.pages:
@@ -39,7 +38,6 @@ def read_pdf(file_obj):
 
 
 def chunk_text(text, chunk_size=800, overlap=200):
-    """Split text into overlapping chunks."""
     words = text.split()
     chunks = []
     start = 0
@@ -54,7 +52,6 @@ def chunk_text(text, chunk_size=800, overlap=200):
 
 
 def build_vectorstore(chunks):
-    """Embed chunks and build FAISS index."""
     if not chunks:
         return None, None
 
@@ -68,7 +65,6 @@ def build_vectorstore(chunks):
 
 
 def retrieve_context(question, state, k=4):
-    """Retrieve top-k relevant chunks."""
     index = state.get("index")
     chunks = state.get("chunks")
 
@@ -83,7 +79,7 @@ def retrieve_context(question, state, k=4):
 
 
 # ============================================================
-# LLM ANSWERING (IBM GRANITE)
+# LLM ANSWERING
 # ============================================================
 
 def generate_answer(question, context):
@@ -100,15 +96,14 @@ def generate_answer(question, context):
             f"Question: {question}\nAnswer:"
         )
 
-    response = llm.generate(
+    response = llm.text_generation(
         prompt,
-        max_new_tokens=512,
+        max_new_tokens=300,
         temperature=0.2,
         top_p=0.9,
-        repetition_penalty=1.05,
     )
 
-    return response.generated_text.strip()
+    return response.strip()
 
 
 # ============================================================
@@ -116,8 +111,6 @@ def generate_answer(question, context):
 # ============================================================
 
 def load_pdf(file, state):
-    """Reset state and load a new PDF."""
-    # Reset everything to avoid memory persistence
     state["chunks"] = None
     state["index"] = None
 
@@ -143,7 +136,6 @@ def load_pdf(file, state):
 
 
 def answer_question(question, state):
-    """Answer a question using the current PDF context."""
     if not question.strip():
         return "Please enter a question."
 
@@ -159,8 +151,8 @@ def answer_question(question, state):
 # GRADIO UI
 # ============================================================
 
-with gr.Blocks(title="Granite PDF QA") as demo:
-    gr.Markdown("## PDF Q&A using IBM Granite\nUpload a PDF and ask questions about its content.")
+with gr.Blocks(title="PDF Q&A Chatbot") as demo:
+    gr.Markdown("## 📄 PDF Q&A Chatbot\nUpload a PDF and ask questions about its content.")
 
     state = gr.State({"chunks": None, "index": None})
 
@@ -168,12 +160,12 @@ with gr.Blocks(title="Granite PDF QA") as demo:
         with gr.Column(scale=1):
             pdf_input = gr.File(label="Upload PDF", file_types=[".pdf"])
             pdf_status = gr.Textbox(label="Status", interactive=False)
-        
+
         with gr.Column(scale=2):
             question = gr.Textbox(label="Your question", lines=2)
             answer = gr.Textbox(label="Answer", lines=8)
             ask_button = gr.Button("Ask")
-    
+
     pdf_input.upload(load_pdf, inputs=[pdf_input, state], outputs=[pdf_status, state])
     ask_button.click(answer_question, inputs=[question, state], outputs=answer)