Update app.py

app.py

@@ -1,8 +1,8 @@
 import gradio as gr
 import os
-import docx
-import numpy as np
+import gc
 import cv2
+import numpy as np
 from pathlib import Path
 from sentence_transformers import SentenceTransformer
 from sklearn.metrics.pairwise import cosine_similarity
@@ -13,50 +13,53 @@ from langchain.memory import ConversationBufferMemory
 from langchain_community.llms import HuggingFaceEndpoint
 from langchain_huggingface import HuggingFaceEmbeddings
 
+# ✅ Semantic model for query validation
+semantic_model = SentenceTransformer("all-MiniLM-L6-v2")
+
+# ✅ Optimized Image Classifier
 class SimpleGeometryClassifier:
     def __init__(self):
         self.reference_embeddings = {
-            "flat.png": {
-                "embedding": None,
-                "label": "Flat or Sheet-Based"
-            },
-            "cylindrical.png": {
-                "embedding": None,
-                "label": "Cylindrical"
-            },
-            "complex.png": {
-                "embedding": None,
-                "label": "Complex Multi Axis Geometry"
-            }
+            "flat.png": {"embedding": None, "label": "Flat or Sheet-Based"},
+            "cylindrical.png": {"embedding": None, "label": "Cylindrical"},
+            "complex.png": {"embedding": None, "label": "Complex Multi Axis Geometry"}
         }
 
     def compute_embedding(self, image_path):
-        img = cv2.imread(image_path)
-        img = cv2.resize(img, (224, 224))
-        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-        
-        win_size = (224, 224)
+        img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
+        img = cv2.resize(img, (128, 128))
+
+        win_size = (128, 128)
         cell_size = (8, 8)
-        block_size = (16, 16)
-        block_stride = (8, 8)
+        block_size = (8, 8)
+        block_stride = (4, 4)
         num_bins = 9
-        
+
         hog = cv2.HOGDescriptor(win_size, block_size, block_stride, cell_size, num_bins)
         embedding = hog.compute(img)
-        
+
+        # ✅ Free OpenCV resources
+        cv2.destroyAllWindows()
+
         return embedding.flatten()
 
     def initialize_reference_embeddings(self, reference_folder):
+        """ Load reference embeddings for classification """
         for image_name in self.reference_embeddings.keys():
             image_path = str(Path(reference_folder) / image_name)
             if Path(image_path).exists():
                 self.reference_embeddings[image_name]["embedding"] = self.compute_embedding(image_path)
             else:
-                print(f"Warning: Reference image {image_path} not found")
+                print(f"Warning: Missing reference image: {image_path}")
+
+    def process_image(self, image_path):
+        """ Classify uploaded image """
+        query_embedding = self.compute_embedding(image_path)
+        return self.find_closest_geometry(query_embedding)
 
     def find_closest_geometry(self, query_embedding):
         best_similarity = -1
-        best_label = None
+        best_label = "Unknown Geometry"
         
         for ref_data in self.reference_embeddings.values():
             if ref_data["embedding"] is not None:
@@ -67,14 +70,16 @@ class SimpleGeometryClassifier:
                     best_similarity = similarity
                     best_label = ref_data["label"]
         
-        return best_label or "Unknown Geometry"
+        return best_label
 
-    def process_image(self, image_path):
-        query_embedding = self.compute_embedding(image_path)
-        return self.find_closest_geometry(query_embedding)
 
-# Initialize semantic model
-semantic_model = SentenceTransformer("all-MiniLM-L6-v2")
+# ✅ Initialize Image Classifier
+classifier = SimpleGeometryClassifier()
+classifier.initialize_reference_embeddings("images")
+
+# ✅ Initialize Chatbot Once
+retriever, qa_chain, embeddings = None, None, None
+retriever, qa_chain, embeddings = initialize_chatbot()
 
 def extract_text_from_docx(file_path):
     doc = docx.Document(file_path)
@@ -154,16 +159,23 @@ def validate_query_semantically(query, retrieved_docs):
     
     return similarity_score >= 0.3
 
-def initialize_chatbot(vector_db, embeddings):
+# ✅ Initialize Chatbot
+def initialize_chatbot():
     memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True, output_key='answer')
-    
+
+    documents = load_documents()
+    embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-base-en-v1.5")
+    vector_db = FAISS.from_documents(documents, embeddings)
+
     retriever = vector_db.as_retriever(search_kwargs={"k": 5})
-    
+
     system_prompt = """You are an AI assistant that answers questions ONLY based on the provided documents.
     - If no relevant documents are retrieved, respond with: "I couldn't find any relevant information."
-    - If the meaning of the query does not match the retrieved documents, say "I couldn't find any relevant information."
-    - Do NOT attempt to answer from general knowledge."""
-    
+    - Do NOT answer from general knowledge."""
+
+    # ✅ Free memory before LLM call
+    gc.collect()
+
     llm = HuggingFaceEndpoint(
         repo_id="mistralai/Mistral-7B-Instruct-v0.3",
         huggingfacehub_api_token=os.environ.get("HUGGINGFACE_API_TOKEN"),
@@ -172,93 +184,57 @@ def initialize_chatbot(vector_db, embeddings):
         task="text-generation",
         system_prompt=system_prompt
     )
-    
+
     qa_chain = ConversationalRetrievalChain.from_llm(
-        llm=llm,
-        retriever=retriever,
-        memory=memory,
-        return_source_documents=True,
-        verbose=False
+        llm=llm, retriever=retriever, memory=memory, return_source_documents=True, verbose=False
     )
-    
-    return retriever, qa_chain
 
-def handle_query(query, history, retriever, qa_chain, embeddings):
+    return retriever, qa_chain, embeddings
+
+def process_image_and_generate_query(image_path):
+    """ Run Image Classification Separately and Generate Query """
+    geometry_type = classifier.process_image(image_path)
+    query = f"I have a {geometry_type} geometry, which screw should I use and what is the best machine to use for {geometry_type} geometry?"
+
+    # ✅ Free up memory **before** calling API
+    gc.collect()
+    return geometry_type, query
+
+def handle_query(query, history):
     retrieved_docs = retrieve_documents(query, retriever, embeddings)
     
-    if not retrieved_docs or not validate_query_semantically(query, retrieved_docs):
+    if not retrieved_docs:
         return history + [(query, "I couldn't find any relevant information.")], ""
     
     response = qa_chain.invoke({"question": query, "chat_history": history})
     assistant_response = response['answer'].strip()
     
-    if not validate_query_semantically(query, retrieved_docs):
-        assistant_response = "I couldn't find any relevant information."
-    
     assistant_response += f"\n\n📄 Source: {', '.join(set(doc.metadata.get('source', 'Unknown') for doc in retrieved_docs))}"
     
     history.append((query, assistant_response))
     return history, ""
 
-def process_image_and_generate_query(image):
-    classifier = SimpleGeometryClassifier()
-    classifier.initialize_reference_embeddings("images")
-    geometry_type = classifier.process_image(image)
-    query = f"I have a {geometry_type} geometry, which screw should I use and what is the best machine to use for {geometry_type} geometry?"
-    return geometry_type, query
 
 def demo():
-    # Initialize classifier
-    classifier = SimpleGeometryClassifier()
-    classifier.initialize_reference_embeddings("images")
-    
-    # Initialize chatbot components
-    documents = load_documents()
-    vector_db, embeddings = create_db(documents)
-    retriever, qa_chain = initialize_chatbot(vector_db, embeddings)
-    
     with gr.Blocks() as app:
-        gr.Markdown("### 🤖 Fastener Agent with Image Recognition 📚")
-        
+        gr.Markdown("### 🔩 Fastener Selection Assistant")
+
         with gr.Row():
             with gr.Column(scale=1):
-                image_input = gr.Image(type="filepath", label="Upload Geometry Image")
-                geometry_label = gr.Textbox(label="Detected Geometry Type", interactive=False)
-                
+                image_input = gr.Image(type="numpy", label="Upload Geometry Image")
+                geometry_label = gr.Textbox(label="Detected Geometry", interactive=False)
+
             with gr.Column(scale=2):
                 chatbot = gr.Chatbot()
-                query_input = gr.Textbox(label="Ask a question about the documents")
+                query_input = gr.Textbox(label="Ask a question")
                 query_btn = gr.Button("Submit")
 
-        def image_upload_handler(image):
-            if image is None:
-                return "", ""
-            geometry_type = classifier.process_image(image)
-            suggested_query = f"I have a {geometry_type} geometry, which screw should I use and what is the best machine to use for {geometry_type} geometry?"
-            return geometry_type, suggested_query
-
-        def user_query_handler(query, history):
-            return handle_query(query, history, retriever, qa_chain, embeddings)
-
-        image_input.change(
-            image_upload_handler,
-            inputs=[image_input],
-            outputs=[geometry_label, query_input]
-        )
-        
-        query_btn.click(
-            user_query_handler,
-            inputs=[query_input, chatbot],
-            outputs=[chatbot, query_input]
-        )
-        
-        query_input.submit(
-            user_query_handler,
-            inputs=[query_input, chatbot],
-            outputs=[chatbot, query_input]
-        )
+        image_input.change(image_upload_handler, inputs=[image_input], outputs=[geometry_label, query_input])
+        query_btn.click(handle_query, inputs=[query_input, chatbot], outputs=[chatbot, query_input])
+        query_input.submit(handle_query, inputs=[query_input, chatbot], outputs=[chatbot, query_input])
 
     app.launch()
 
+
 if __name__ == "__main__":
     demo()