Update app.py

app.py

@@ -2,7 +2,8 @@ import gradio as gr
 from typing import List, Dict, Tuple
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, AutoModel
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+from sentence_transformers import SentenceTransformer
 import torch
 import os
 from astrapy.db import AstraDB
@@ -10,8 +11,8 @@ from dotenv import load_dotenv
 from huggingface_hub import login
 import time
 import logging
-from functools import lru_cache
 import numpy as np
+from functools import lru_cache
 
 # Configure logging
 logging.basicConfig(
@@ -34,7 +35,7 @@ class LegalTextSearchBot:
             )
             self.collection = self.astra_db.collection(os.getenv("ASTRA_DB_COLLECTION"))
             
-            # Initialize language model for text generation
+            # Initialize language model
             model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
             model = AutoModelForCausalLM.from_pretrained(
                 model_name,
@@ -56,13 +57,8 @@ class LegalTextSearchBot:
             )
             self.llm = HuggingFacePipeline(pipeline=pipe)
             
-            # Initialize embedding model
-            self.embedding_model_name = "sentence-transformers/all-MiniLM-L6-v2"
-            self.embedding_pipeline = pipeline(
-                "feature-extraction",
-                model=self.embedding_model_name,
-                device_map="auto"
-            )
+            # Initialize sentence transformer for embeddings
+            self.embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
             
             self.template = """
             IMPORTANT: You are a legal assistant that provides accurate information based on the Indian legal sections provided in the context.
@@ -87,6 +83,8 @@ class LegalTextSearchBot:
             self.chat_history = ""
             self.is_searching = False
             
+            logger.info("Successfully initialized LegalTextSearchBot")
+            
         except Exception as e:
             logger.error(f"Error initializing LegalTextSearchBot: {str(e)}")
             raise
@@ -96,21 +94,19 @@ class LegalTextSearchBot:
         try:
             # Clean and prepare text
             text = text.replace('\n', ' ').strip()
+            if not text:
+                text = " "  # Ensure non-empty input
             
             # Generate embedding
-            outputs = self.embedding_pipeline(text)
-            embeddings = torch.mean(torch.tensor(outputs[0]), dim=0)
-            
-            # Convert to list and ensure correct dimension
-            embedding_list = embeddings.tolist()
+            embedding = self.embedding_model.encode(text)
             
-            # Pad or truncate to exactly 1024 dimensions
-            if len(embedding_list) < 1024:
-                embedding_list.extend([0.0] * (1024 - len(embedding_list)))
-            elif len(embedding_list) > 1024:
-                embedding_list = embedding_list[:1024]
+            # Pad or truncate to 1024 dimensions
+            if len(embedding) < 1024:
+                embedding = np.pad(embedding, (0, 1024 - len(embedding)))
+            elif len(embedding) > 1024:
+                embedding = embedding[:1024]
             
-            return embedding_list
+            return embedding.tolist()
             
         except Exception as e:
             logger.error(f"Error generating embedding: {str(e)}")
@@ -125,7 +121,7 @@ class LegalTextSearchBot:
             
             results = list(self.collection.vector_find(
                 query_embedding,
-                limit=5,
+                top_k=5,  # Using top_k instead of limit
                 fields=["section_number", "title", "chapter_number", "chapter_title", 
                        "content", "type", "metadata"]
             ))
@@ -142,10 +138,13 @@ class LegalTextSearchBot:
             results = list(self._cached_search(query))
             
             if not results and self.is_searching:
-                results = list(self.collection.find(
-                    {},
-                    limit=5
-                ))
+                # Fallback to regular search
+                cursor = self.collection.find({})
+                results = []
+                for doc in cursor:
+                    if len(results) >= 5:
+                        break
+                    results.append(doc)
             
             return results