Update app.py

app.py

@@ -1,12 +1,12 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from peft import PeftModel
-from langchain.text_splitter import CharacterTextSplitter
-from langchain.docstore.document import Document
 from langchain_community.llms import HuggingFacePipeline
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
-from langchain.chains.retrieval_qa.base import RetrievalQA
+from langchain.text_splitter import CharacterTextSplitter
+from langchain.docstore.document import Document
+from langchain.chains import RetrievalQA
 
 # Model and Tokenizer
 model_name = "Meldashti/chatbot"
@@ -15,55 +15,64 @@ tokenizer = AutoTokenizer.from_pretrained("unsloth/Llama-3.2-3B")
 
 # Merge PEFT weights with base model
 model = PeftModel.from_pretrained(base_model, model_name)
-model = model.merge_and_unload()  # This merges the PEFT weights into the base model
+model = model.merge_and_unload()
 
-# Set up the text-generation pipeline
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+# Simplified pipeline with minimal parameters
+generator = pipeline(
+    "text-generation", 
+    model=model, 
+    tokenizer=tokenizer,
+    max_new_tokens=50,  # Very low to test responsiveness
+    do_sample=False
+)
 
-# Use the HuggingFacePipeline from langchain_community
+# LLM wrapper
 llm = HuggingFacePipeline(pipeline=generator)
 
-# Initialize Hugging Face embeddings model
-embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
+# Embeddings
+embeddings = HuggingFaceEmbeddings(model_name="paraphrase-MiniLM-L3-v2")
 
-# Sample documents
+# Sample documents (minimal)
 documents = [
-    Document(page_content="Document 1 content goes here..."),
-    Document(page_content="Document 2 content goes here..."),
-    # Add more documents as needed
+    Document(page_content="Example document about food industry caps"),
+    Document(page_content="Information about manufacturing processes")
 ]
 
-# Split documents into smaller chunks for better retrieval
-text_splitter = CharacterTextSplitter(chunk_size=500, chunk_overlap=100)
+# Text splitting
+text_splitter = CharacterTextSplitter(chunk_size=100, chunk_overlap=20)
 split_documents = text_splitter.split_documents(documents)
 
-# Create FAISS vector store
+# Vector store
 vector_store = FAISS.from_documents(split_documents, embeddings)
+retriever = vector_store.as_retriever(search_kwargs={"k": 2})
 
-# Define a retriever that uses FAISS vector store
-retriever = vector_store.as_retriever()
-
-# Create Retrieval QA Chain
+# Retrieval QA Chain
 rag_chain = RetrievalQA.from_chain_type(
-    llm=llm,
-    chain_type="stuff",
+    llm=llm, 
+    chain_type="stuff", 
     retriever=retriever
 )
 
-# Define the chat function
+# Chat function with extensive logging
 def chat(message, history):
     print(f"Received message: {message}")
     try:
-        response = rag_chain.invoke({"query": message})
-        print(f"Response generated: {response}")
-        return str(response['result'])
+        # Add timeout mechanism
+        import timeout_decorator
+
+        @timeout_decorator.timeout(10)  # 10 seconds timeout
+        def get_response():
+            response = rag_chain.invoke({"query": message})
+            return str(response['result'])
+
+        return get_response()
     except Exception as e:
-        print(f"Error generating response: {e}")
-        return "Sorry, I couldn't generate a response."
+        print(f"Error generating response: {type(e)}, {e}")
+        return f"An error occurred: {str(e)}"
 
-# Set up the Gradio interface
+# Gradio interface
 demo = gr.ChatInterface(chat, type="messages", autofocus=False)
 
-# Launch the app
+# Launch
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(debug=True)