Update app.py

app.py

@@ -66,6 +66,9 @@ def load_document(file: NamedTemporaryFile, parser: str = "llamaparse") -> List[
     else:
         raise ValueError("Invalid parser specified. Use 'pypdf' or 'llamaparse'.")
 
+import requests
+import numpy as np
+
 class HuggingFaceEmbeddings:
     def __init__(self, api_token):
         self.api_url = "https://api-inference.huggingface.co/models/dunzhang/stella_en_1.5B_v5"
@@ -76,12 +79,44 @@ class HuggingFaceEmbeddings:
         return response.json()
 
     def embed_documents(self, texts):
-        payload = {"inputs": texts}
-        response = self.query(payload)
-        if isinstance(response, list):
-            return [np.array(embedding) for embedding in response]
-        else:
-            raise ValueError(f"Unexpected response format: {response}")
+        # Split long texts into smaller chunks
+        max_chunk_length = 512  # Adjust this value based on the model's requirements
+        chunked_texts = []
+        for text in texts:
+            if len(text) > max_chunk_length:
+                chunks = [text[i:i+max_chunk_length] for i in range(0, len(text), max_chunk_length)]
+                chunked_texts.extend(chunks)
+            else:
+                chunked_texts.append(text)
+
+        # Process chunks in batches
+        batch_size = 8  # Adjust this value based on API limits and performance
+        all_embeddings = []
+        for i in range(0, len(chunked_texts), batch_size):
+            batch = chunked_texts[i:i+batch_size]
+            payload = {
+                "inputs": batch,
+                "task": "sentence-similarity"  # Specify the task
+            }
+            response = self.query(payload)
+            if isinstance(response, list):
+                all_embeddings.extend(response)
+            elif isinstance(response, dict) and 'error' in response:
+                raise ValueError(f"API Error: {response['error']}")
+            else:
+                raise ValueError(f"Unexpected response format: {response}")
+
+        # Average embeddings for chunks of the same original text
+        final_embeddings = []
+        i = 0
+        for text in texts:
+            num_chunks = max(1, len(text) // max_chunk_length)
+            text_embeddings = all_embeddings[i:i+num_chunks]
+            avg_embedding = np.mean(text_embeddings, axis=0)
+            final_embeddings.append(avg_embedding)
+            i += num_chunks
+
+        return final_embeddings
 
     def embed_query(self, text):
         return self.embed_documents([text])[0]
@@ -114,7 +149,12 @@ def update_vectors(files, parser):
         logging.warning("No files provided for update_vectors")
         return "Please upload at least one PDF file.", display_documents()
     
-    embed = get_embeddings()
+    try:
+        embed = get_embeddings()
+    except Exception as e:
+        logging.error(f"Error initializing embeddings: {str(e)}")
+        return f"Error initializing embeddings: {str(e)}", display_documents()
+
     total_chunks = 0
     
     all_data = []
@@ -163,7 +203,6 @@ def update_vectors(files, parser):
     save_documents(uploaded_documents)
 
     return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files using {parser}.", display_documents()
-
 def delete_documents(selected_docs):
     global uploaded_documents