mxbai endpoint

.gitignore

@@ -0,0 +1,31 @@
+lib
+dist
+allWatches.json
+vectors.json
+node_modules/
+/test-results/
+/playwright-report/
+/blob-report/
+/playwright/.cache/
+
+
+src/all_shopping_scrape_results/*
+src/all_product_scrape_results/*
+.envrc
+tsconfig.tsbuildinfo
+.tshy
+.tshy-build
+response-cache
+venv
+*.html
+response2.json
+code.py
+mic-scrape.json
+*.json
+*.parquet
+__pycache__
+.DS_Store
+.passwd-s3fs
+.idea/*
+myenv/
+env/

app.py

@@ -2,36 +2,32 @@ import torch
 import src.constants.config as configurations
 from sentence_transformers import SentenceTransformer
 from sentence_transformers import CrossEncoder
-from src.constants.credentials import cohere_trial_key
+from src.constants.credentials import cohere_trial_key, mixedbread_key
 import streamlit as st
 from src.reader import Reader
 from src.utils_search import UtilsSearch
 from copy import deepcopy
 import numpy as np
 import cohere
+from mixedbread_ai.client import MixedbreadAI
+from src.pytorch_modules.datasets.schema_string_dataset import SchemaStringDataset
 
 
 configurations = configurations.service_mxbai_msc_direct_config
 api_key = cohere_trial_key
 co = cohere.Client(api_key)
 semantic_column_names = configurations["semantic_column_names"]
-
-# Check CUDA availability and set device
-if torch.cuda.is_available():
-    torch.cuda.set_device(0)  # Use the first GPU
-else:
-    st.write("CUDA is not available. Using CPU instead.")
+model = MixedbreadAI(api_key=mixedbread_key)
+cross_encoder_name = configurations["cross_encoder_name"]
 
 @st.cache_data
 def init():
     config = configurations
     search_utils = UtilsSearch(config)
     reader = Reader(config=config["reader_config"])
-    model = SentenceTransformer(config['sentence_transformer_name'], device='cuda:0')
-    cross_encoder = CrossEncoder(config['cross_encoder_name'], device='cuda:0')
     df = reader.read()
     index = search_utils.dataframe_to_index(df)
-    return df, model, cross_encoder, index, search_utils
+    return df, index, search_utils
 
 def get_possible_values_for_column(column_name, search_utils, df):
     if column_name not in st.session_state:
@@ -44,14 +40,15 @@ if 'init_results' not in st.session_state:
     st.session_state.init_results = init()
 
 # Now you can access your initialized objects directly from the session state
-df, model, cross_encoder, index, search_utils = st.session_state.init_results
+df, index, search_utils = st.session_state.init_results
 
 # Streamlit app layout
 st.title('Search Demo')
 
 # Input fields
 query = st.text_input('Enter your search query here')
-use_cohere = st.checkbox('Use Cohere', value=False)  # Default to checked
+# use_cohere = st.checkbox('Use Cohere', value=False)  # Default to checked
+use_cohere = False
 
 programmatic_search_config = deepcopy(configurations['programmatic_search_config'])
 
@@ -87,21 +84,31 @@ programmatic_search_config['discrete_columns'] = dynamic_programmatic_search_con
 # Search button
 if st.button('Search'):
     if query:  # Checking if a query was entered
-        df_filtered = search_utils.filter_dataframe(df, programmatic_search_config)
+        df_retrieved = search_utils.retrieve(query, df, model, index, top_k=1000, api=True)
+        df_filtered = search_utils.filter_dataframe(df_retrieved, programmatic_search_config)
+        df_filtered = df_filtered.sort_values(by='similarities', ascending=True)
+        df_filtered = df_filtered[:100].reset_index(drop=True)
+
         if len(df_filtered) == 0:
             st.write('No results found')
         else:
-            index = search_utils.dataframe_to_index(df_filtered)
             if use_cohere == False:
-                # Call your Cohere-based search function here
-                results_df = search_utils.search(query, df_filtered, model, cross_encoder, index)
-                results_df = search_utils.drop_columns(results_df, programmatic_search_config)
+                records = df_filtered.to_dict(orient='records')
+                dataset_str = SchemaStringDataset(records, configurations)
+                documents = [batch["inputs"][:256] for batch in dataset_str]
+                res = model.reranking(
+                    model=cross_encoder_name,
+                    query=query,
+                    input=documents,
+                    top_k=10,
+                    return_input=False
+                )
+                ids = [item.index for item in res.data]
+                results_df = df_filtered.loc[ids]
 
             else:
-                df_retrieved = search_utils.retrieve(query, df_filtered, model, index)
-                df_retrieved = search_utils.drop_columns(df_retrieved, programmatic_search_config)
-                df_retrieved.fillna(value="", inplace=True)
-                docs = df_retrieved.to_dict('records')
+                df_filtered.fillna(value="", inplace=True)
+                docs = df_filtered.to_dict('records')
                 column_names = semantic_column_names
                 docs = [{name: str(doc[name]) for name in column_names} for doc in docs]
                 rank_fields = list(docs[0].keys())
@@ -109,7 +116,7 @@ if st.button('Search'):
                                     rank_fields=rank_fields)
                 top_ids = [hit.index for hit in results.results]
                 # Create the DataFrame with the rerank results
-                results_df = df_retrieved.iloc[top_ids].copy()
+                results_df = df_filtered.iloc[top_ids].copy()
                 results_df['rank'] = (np.arange(len(results_df)) + 1)
 
             st.write(results_df)

requirements.txt

@@ -8,4 +8,5 @@ s3fs
 numpy
 faiss-gpu
 sentence_transformers
-cohere
+cohere
+mixedbread_ai

src/constants/config.py

@@ -1,52 +1,6 @@
 import src.constants.credentials as cred
 import os
 
-service_mxbai_made_in_china_config = {"reader_config": {"input_path": os.environ['made_in_china_s3_path'],
-                                                                     "credentials": cred.credentials_backblaze,
-                                                                     "format":"parquet"
-                                                                    },
-                             "sample_size": 32,
-                             "sentence_transformer_name": "mixedbread-ai/mxbai-embed-large-v1",
-                             "cross_encoder_name": "mixedbread-ai/mxbai-rerank-large-v1",
-                             "batch_size": 4,
-                             "dataset_size": 32,
-                             "seq_len": 256,
-                             "top_k": 1000,
-                             "programmatic_search_config": {
-                                 "scalar_columns": [{"column_name": "price", "min_value": 0, "max_value": "10000"}],
-                                 "discrete_columns": [{"column_name": "supplierName",
-                                                       # "default_values": ['Zhongshan Norye Hardware Co., Ltd.']
-                                                       "default_values": []
-                                                       },
-                                                      {"column_name": "warranty",
-                                                       # "default_values": ['Zhongshan Norye Hardware Co., Ltd.']
-                                                       "default_values": []
-                                                       }
-                                                      ],
-                                 "columns_to_drop": ["similarities", "embeddings"]
-                             }
-                             }
-
-
-service_mxbai_msc_direct_sample_config = {"reader_config": {"input_path": os.environ['msc_direct_s3_path'],
-                                                                     "credentials": cred.credentials_backblaze,
-                                                                     "format":"parquet"
-                                                                    },
-                             "sample_size": 32,
-                             "sentence_transformer_name": "mixedbread-ai/mxbai-embed-large-v1",
-                             "cross_encoder_name": "mixedbread-ai/mxbai-rerank-large-v1",
-                             "batch_size": 4,
-                             "dataset_size": 32,
-                             "seq_len": 256,
-                             "top_k": 50,
-                            "semantic_column_names": ['name', 'price', 'brand', 'keyword', 'description',
-                                                     'specifications'],
-                             "programmatic_search_config": {
-                                 "scalar_columns": [{"column_name": "price", "min_value": 0, "max_value": "10000"}],
-                                 "discrete_columns": [{"column_name": "brand", "default_values": []}],
-                                 "columns_to_drop": ["similarities", "embeddings", "index"]
-                             }
-                             }
 
 service_mxbai_msc_direct_config = {"reader_config": {"input_path": os.environ['msc_direct_s3_path'],
                                                                      "credentials": cred.credentials_backblaze,

src/constants/credentials.py

@@ -8,4 +8,5 @@ credentials_backblaze = {"access_key_id": os.environ['credentials_backblaze_acce
                          }
 
 
-cohere_trial_key = os.environ["cohere_trial_key"]
+cohere_trial_key = os.environ["cohere_trial_key"]
+mixedbread_key = os.environ["mixedbread_key"]

src/utils_search.py

@@ -21,8 +21,8 @@ class UtilsSearch:
         index.add(norm_embeddings)
         return index  # Ad
 
-    @staticmethod
-    def retrieve(query, df, model, index, top_k=100):
+
+    def retrieve(self, query, df, model, index, top_k=100, api=False):
         query += "Represent this sentence for searching relevant passages: "
         """
         Search the DataFrame for the given query and return a sorted DataFrame based on similarity.
@@ -35,14 +35,24 @@ class UtilsSearch:
         :return: A new DataFrame sorted by similarity to the query, with a 'similarities' column.
         """
         # Check if CUDA is available and set the device accordingly
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        model.to(device)
+        if not api:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            model.to(device)
+
+            # Compute the query embedding
+            query_vector = model.encode(query, convert_to_tensor=True, device=device).cpu().numpy()
+            query_vector /= np.linalg.norm(query_vector)
+        else:
+            res = model.embeddings(
+                input=[query],
+                model=self.config["sentence_transformer_name"],
+                prompt=None,
+            )
+            query_vector = np.array([entry.embedding for entry in res.data][0]).astype(np.float32)
 
-        # Compute the query embedding
-        query_vector = model.encode(query, convert_to_tensor=True, device=device).cpu().numpy()
 
         # Normalize the query vector
-        query_vector /= np.linalg.norm(query_vector)
+
 
         # Perform the search
         distances, indices = index.search(np.array([query_vector]), top_k)
@@ -55,7 +65,7 @@ class UtilsSearch:
         retrieved_df = retrieved_df.assign(similarities=distances[0])
 
         if 'similarities' in retrieved_df.columns:
-            retrieved_df = retrieved_df.sort_values(by='similarities', ascending=False)
+            retrieved_df = retrieved_df.sort_values(by='similarities', ascending=True)
 
         # Optionally, you might want to reset the index if the order matters or if you need to serialize the DataFrame without index issues
         retrieved_df = retrieved_df.reset_index(drop=True)
@@ -149,5 +159,3 @@ class UtilsSearch:
         columns_to_drop = config.get('columns_to_drop', [])
         df_dropped = df.drop(columns_to_drop, axis=1)
         return df_dropped
-
-