Refactor embedding model integration and update API documentation for search response format

[all_embedded] The Alchemy of Happiness (Ghazzālī, Claud Field) (Z-Library).parquet → [embed] The Alchemy of Happiness (Ghazzālī, Claud Field).parquet

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ced650f23166f55939fb6dfec6df2fd7d83995a9db362a1a7460d36e6f3ab510
-size 3118786
+oid sha256:ca01a279b52f21c7e7d8441f8145f20201a255d8c015f3059920b1b957726a61
+size 4232361

main.py

@@ -51,7 +51,7 @@ class QueryInput(BaseModel):
 class SearchResult(BaseModel):
     text: str
     similarity: float
-    model_type: str
+    model_type: Literal["WhereIsAI_UAE_Large_V1", "BAAI_bge_large_en_v1.5"] 
 
 class TokenResponse(BaseModel):
     access_token: str
@@ -73,10 +73,13 @@ class RefreshRequest(BaseModel):
     refresh_token: str
     
 # Cache management
-@lru_cache(maxsize=1)
-def get_sentence_transformer():
-    """Load and cache the SentenceTransformer model with lru_cache"""
-    return SentenceTransformer(model_name_or_path="all-mpnet-base-v2", device="cpu")
+@lru_cache(maxsize=2)  # Cache both models
+def get_embedding_models():
+    """Load and cache both embedding models"""
+    return {
+        "uae-large": SentenceTransformer("WhereIsAI/UAE-Large-V1", device="cpu"),
+        "bge-large": SentenceTransformer("BAAI/bge-large-en-v1.5", device="cpu")
+    }
 
 def get_cached_embeddings(text: str, model_type: str) -> Optional[List[float]]:
     """Try to get embeddings from cache"""
@@ -91,7 +94,7 @@ def set_cached_embeddings(text: str, model_type: str, embeddings: List[float]):
 @lru_cache(maxsize=1)
 def load_dataframe():
     """Load and cache the parquet dataframe"""
-    database_file = Path(__file__).parent / "[all_embedded] The Alchemy of Happiness (Ghazzālī, Claud Field) (Z-Library).parquet"
+    database_file = Path(__file__).parent / "[embed] The Alchemy of Happiness (Ghazzālī, Claud Field).parquet"
     return pd.read_parquet(database_file)
 
 # Utility functions
@@ -102,61 +105,53 @@ def cosine_similarity(embedding_0, embedding_1):
     return dot_product / (norm_0 * norm_1)
 
 def generate_embedding(model, text: str, model_type: str) -> List[float]:
-    # Try to get from cache first
     cached_embedding = get_cached_embeddings(text, model_type)
     if cached_embedding is not None:
         return cached_embedding
 
-    # Generate new embedding if not in cache
-    if model_type == "all-mpnet-base-v2":
-        chunk_embedding = model.encode(
-            text,
-            convert_to_tensor=True
-        )
-        embedding = np.array(t.Tensor.cpu(chunk_embedding)).tolist()
-    elif model_type == "text-embedding-3-small":
-        response = model.embeddings.create(
-            input=text,
-            model="text-embedding-3-small"
-        )
-        embedding = response.data[0].embedding
-
-    # Cache the new embedding
+    # Generate new embedding
+    embedding = model.encode(
+        text,
+        convert_to_tensor=True,
+        normalize_embeddings=True  # Important for UAE and BGE models
+    )
+    embedding = np.array(t.Tensor.cpu(embedding)).tolist()
+    
     set_cached_embeddings(text, model_type, embedding)
     return embedding
 
-def search_query(client, st_model, query: str, df: pd.DataFrame, n: int = 1) -> List[Dict]:
-    # Generate embeddings for both models
-    mpnet_embedding = generate_embedding(st_model, query, "all-mpnet-base-v2")
-    openai_embedding = generate_embedding(client, query, "text-embedding-3-small")
+def search_query(st_models, query: str, df: pd.DataFrame, n: int = 1) -> List[Dict]:
+    # Generate embeddings with both models
+    uae_embedding = generate_embedding(st_models["uae-large"], query, "uae-large")
+    bge_embedding = generate_embedding(st_models["bge-large"], query, "bge-large")
     
     # Calculate similarities
-    df['mpnet_similarities'] = df.all_mpnet_embedding.apply(
-        lambda x: cosine_similarity(x, mpnet_embedding)
+    df['uae_similarities'] = df["WhereIsAI_UAE_Large_V1"].apply(
+        lambda x: cosine_similarity(x, uae_embedding)
     )
-    df['openai_similarities'] = df.openai_embedding.apply(
-        lambda x: cosine_similarity(x, openai_embedding)
+    df['bge_similarities'] = df["BAAI_bge_large_en_v1.5"].apply(
+        lambda x: cosine_similarity(x, bge_embedding)
     )
     
     # Get top results for each model
-    mpnet_results = df.nlargest(n, 'mpnet_similarities')
-    openai_results = df.nlargest(n, 'openai_similarities')
+    uae_results = df.nlargest(n, 'uae_similarities')
+    bge_results = df.nlargest(n, 'bge_similarities')
     
     # Format results
     results = []
     
-    for _, row in mpnet_results.iterrows():
+    for _, row in uae_results.iterrows():
         results.append({
             "text": row["ext"],
-            "similarity": float(row["mpnet_similarities"]),
-            "model_type": "all-mpnet-base-v2"
+            "similarity": float(row["uae_similarities"]),
+            "model_type": "WhereIsAI_UAE_Large_V1"
         })
     
-    for _, row in openai_results.iterrows():
+    for _, row in bge_results.iterrows():
         results.append({
             "text": row["ext"],
-            "similarity": float(row["openai_similarities"]),
-            "model_type": "text-embedding-3-small"
+            "similarity": float(row["bge_similarities"]),
+            "model_type": "BAAI_bge_large_en_v1.5"
         })
     
     return results
@@ -309,17 +304,14 @@ def logout(
 
 @app.post("/search", response_model=List[SearchResult])
 async def search(
-    query_input: QueryInput,
-    username: str = Depends(verify_access_token),
-):
+            query_input: QueryInput,
+            username: str = Depends(verify_access_token),
+        ):
     try:
-        # Initialize clients using cached functions
-        client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
-        st_model = get_sentence_transformer()
+        st_models = get_embedding_models()
         df = load_dataframe()
         
-        # Perform search with both models
-        results = search_query(client, st_model, query_input.query, df, n=1)
+        results = search_query(st_models, query_input.query, df, n=1)
         return [SearchResult(**result) for result in results]
     
     except Exception as e:

static/index.html

@@ -138,18 +138,18 @@ class LoginResponse {
                   <p>This endpoint is used to send a search query and retrieve results. It requires a valid access token.</p>
                   
                   <h4>Response:</h4>
-                  <pre><code class="language-json">[
-                  {
-                      "text": "Result 1 text",
-                      "similarity": 0.95,
-                      "model_type": "all-mpnet-base-v2"
-                  },
-                  {
-                      "text": "Result 2 text",
-                      "similarity": 0.92,
-                      "model_type": "text-embedding-3-small"
-                  }
-              ]</code></pre>
+                <pre><code class="language-json">[
+                    {
+                        "text": "Result 1 text",
+                        "similarity": 0.95,
+                        "model_type": "UAE-Large-V1"
+                    },
+                    {
+                        "text": "Result 2 text",
+                        "similarity": 0.92,
+                        "model_type": "BGE-Large-V1.5"
+                    }
+                ]</code></pre>
               </div>
 
                 <div class="endpoint" id="save">