Update app.py

app.py

@@ -18,6 +18,10 @@ from huggingface_hub import InferenceClient
 import inspect
 import logging
 import shutil
+import time
+from tenacity import retry, stop_after_attempt, wait_exponential
+import requests
+import numpy as np
 
 
 # Set up basic configuration for logging
@@ -66,21 +70,21 @@ 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"
         self.headers = {"Authorization": f"Bearer {api_token}"}
 
+    @retry(stop=stop_after_attempt(5), wait=wait_exponential(multiplier=1, min=4, max=10))
     def query(self, payload):
         response = requests.post(self.api_url, headers=self.headers, json=payload)
-        return response.json()
+        result = response.json()
+        if 'error' in result and 'is currently loading' in result['error']:
+            raise Exception("Model is still loading")
+        return result
 
     def embed_documents(self, texts):
-        # Split long texts into smaller chunks
-        max_chunk_length = 512  # Adjust this value based on the model's requirements
+        max_chunk_length = 512
         chunked_texts = []
         for text in texts:
             if len(text) > max_chunk_length:
@@ -89,14 +93,13 @@ class HuggingFaceEmbeddings:
             else:
                 chunked_texts.append(text)
 
-        # Process chunks in batches
-        batch_size = 8  # Adjust this value based on API limits and performance
+        batch_size = 8
         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
+                "task": "sentence-similarity"
             }
             response = self.query(payload)
             if isinstance(response, list):
@@ -106,7 +109,6 @@ class HuggingFaceEmbeddings:
             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:
@@ -203,6 +205,7 @@ 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