Update app.py

app.py

@@ -2,7 +2,7 @@ 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
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, AutoModel
 import torch
 import os
 from astrapy.db import AstraDB
@@ -11,6 +11,7 @@ from huggingface_hub import login
 import time
 import logging
 from functools import lru_cache
+import numpy as np
 
 # Configure logging
 logging.basicConfig(
@@ -23,25 +24,26 @@ logger = logging.getLogger(__name__)
 load_dotenv()
 login(token=os.getenv("HUGGINGFACE_API_TOKEN"))
 
-# Initialize model with CPU-compatible settings
-model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    device_map="auto",
-    torch_dtype=torch.float32,  # Use float32 for CPU compatibility
-)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-
 class LegalTextSearchBot:
     def __init__(self):
         try:
+            # Initialize AstraDB connection
             self.astra_db = AstraDB(
                 token=os.getenv("ASTRA_DB_APPLICATION_TOKEN"),
                 api_endpoint=os.getenv("ASTRA_DB_API_ENDPOINT")
             )
-            self.collection = self.astra_db.collection("legal_content")
+            self.collection = self.astra_db.collection(os.getenv("ASTRA_DB_COLLECTION"))
             
-            # Initialize pipeline with CPU settings
+            # Initialize language model for text generation
+            model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                device_map="auto",
+                torch_dtype=torch.float32,
+            )
+            tokenizer = AutoTokenizer.from_pretrained(model_name)
+            
+            # Initialize text generation pipeline
             pipe = pipeline(
                 "text-generation",
                 model=model,
@@ -54,6 +56,14 @@ 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"
+            )
+            
             self.template = """
             IMPORTANT: You are a legal assistant that provides accurate information based on the Indian legal sections provided in the context.
             
@@ -81,17 +91,45 @@ class LegalTextSearchBot:
             logger.error(f"Error initializing LegalTextSearchBot: {str(e)}")
             raise
 
+    def get_embedding(self, text: str) -> List[float]:
+        """Generate embedding vector for text"""
+        try:
+            # Clean and prepare text
+            text = text.replace('\n', ' ').strip()
+            
+            # 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()
+            
+            # 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]
+            
+            return embedding_list
+            
+        except Exception as e:
+            logger.error(f"Error generating embedding: {str(e)}")
+            raise
+
     @lru_cache(maxsize=100)
     def _cached_search(self, query: str) -> tuple:
-        """Cached version of vector search to improve performance"""
+        """Cached version of vector search"""
         try:
+            # Generate embedding for query
+            query_embedding = self.get_embedding(query)
+            
             results = list(self.collection.vector_find(
-                query,
+                query_embedding,
                 limit=5,
                 fields=["section_number", "title", "chapter_number", "chapter_title", 
                        "content", "type", "metadata"]
             ))
-            return tuple(results)  # Convert to tuple for caching
+            return tuple(results)
         except Exception as e:
             logger.error(f"Error in vector search: {str(e)}")
             return tuple()