Update add_embeddings.py

add_embeddings.py

@@ -1,19 +1,34 @@
 import os
 from PyPDF2 import PdfReader
-from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModel
+import torch
 import chromadb
 from typing import List, Dict
 import re
+import numpy as np
 
 class LegalDocumentProcessor:
     def __init__(self):
-        self.embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+        self.tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
+        self.model = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
         self.chroma_client = chromadb.Client()
         self.collection = self.chroma_client.create_collection(
             name="indian_legal_docs",
             metadata={"description": "Indian Criminal Law Documents"}
         )
         
+    def mean_pooling(self, model_output, attention_mask):
+        token_embeddings = model_output[0]
+        input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+        return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
+        
+    def get_embedding(self, text: str) -> List[float]:
+        inputs = self.tokenizer(text, padding=True, truncation=True, max_length=512, return_tensors='pt')
+        with torch.no_grad():
+            model_output = self.model(**inputs)
+        sentence_embeddings = self.mean_pooling(model_output, inputs['attention_mask'])
+        return sentence_embeddings[0].tolist()
+        
     def process_pdf(self, pdf_path: str) -> List[str]:
         """Extract text from PDF and split into chunks"""
         reader = PdfReader(pdf_path)
@@ -21,13 +36,11 @@ class LegalDocumentProcessor:
         for page in reader.pages:
             text += page.extract_text()
             
-        # Split into meaningful chunks (by sections/paragraphs)
         chunks = self._split_into_chunks(text)
         return chunks
     
     def _split_into_chunks(self, text: str, max_chunk_size: int = 1000) -> List[str]:
         """Split text into smaller chunks while preserving context"""
-        # Split on section boundaries or paragraphs
         sections = re.split(r'(Chapter \d+|Section \d+|\n\n)', text)
         
         chunks = []
@@ -55,16 +68,14 @@ class LegalDocumentProcessor:
         }
         
         for law_code, pdf_path in pdf_files.items():
-            # Process PDF
             chunks = self.process_pdf(pdf_path)
             
-            # Generate embeddings and store in ChromaDB
             for i, chunk in enumerate(chunks):
-                embeddings = self.embedding_model.encode([chunk]).tolist()
+                embedding = self.get_embedding(chunk)
                 
                 self.collection.add(
                     documents=[chunk],
-                    embeddings=embeddings,
+                    embeddings=[embedding],
                     metadatas=[{
                         "law_code": law_code,
                         "chunk_id": f"{law_code}_chunk_{i}",
@@ -73,11 +84,11 @@ class LegalDocumentProcessor:
                     ids=[f"{law_code}_chunk_{i}"]
                 )
                 
-    def search_documents(self, query: str, n_results: int = 3) -> List[Dict]:
+    def search_documents(self, query: str, n_results: int = 3) -> Dict:
         """Search for relevant legal information"""
-        query_embedding = self.embedding_model.encode([query]).tolist()
+        query_embedding = self.get_embedding(query)
         results = self.collection.query(
-            query_embeddings=query_embedding,
+            query_embeddings=[query_embedding],
             n_results=n_results
         )
         
@@ -87,11 +98,9 @@ class LegalDocumentProcessor:
         }
 
 if __name__ == "__main__":
-    # Initialize and run document processing
     processor = LegalDocumentProcessor()
     processor.process_and_store_documents()
     
-    # Test search functionality
     test_query = "What are the provisions for digital evidence?"
     results = processor.search_documents(test_query)
     print(f"Query: {test_query}")