ug

Author: bstadt

FlagEmbedding/inference/embedder/encoder_only/voyage.py

@@ -39,6 +39,75 @@ class VoyageMockModel():
     def __init__(self):
         self.config = VoyageEmbedderConfig()
 
+MODEL = 'voyage-3-large'
+
+def call_api_query(text_chunk, cache_dir="./voyage_cache"):
+    """Simple function that just calls the API - used in multiprocessing"""
+    cache_dir = pathlib.Path(cache_dir)
+    cache_dir.mkdir(exist_ok=True)
+    
+    # Create cache key from input
+    chunks_str = json.dumps(text_chunk, sort_keys=True)
+    cache_key = hashlib.md5(chunks_str.encode()).hexdigest()
+    cache_file = cache_dir / f"query_{cache_key}.npy"
+    
+    # Check cache first
+    if cache_file.exists():
+        try:
+            return np.load(cache_file).tolist()
+        except Exception as e:
+            print(f"Failed to load cache file {cache_file}: {e}")
+    
+    # If not in cache, call API
+    vo = voyageai.Client(api_key=VOYAGE_API_KEY)
+    max_retries = 5
+    for attempt in range(max_retries):
+        try:
+            result = vo.embed(text_chunk, model=MODEL, input_type='query')
+            embeddings = result.embeddings
+            # Save to cache
+            np.save(cache_file, embeddings)
+            return embeddings
+        except Exception as e:
+            if attempt == max_retries - 1:
+                print(f"Failed after {max_retries} attempts: {str(e)}")
+                raise e
+            time.sleep(5)
+
+def call_api_document(text_chunk, cache_dir="./voyage_cache"):
+    """Simple function that just calls the API - used in multiprocessing"""
+    cache_dir = pathlib.Path(cache_dir)
+    cache_dir.mkdir(exist_ok=True)
+    
+    # Create cache key from input
+    chunks_str = json.dumps(text_chunk, sort_keys=True)
+    cache_key = hashlib.md5(chunks_str.encode()).hexdigest()
+    cache_file = cache_dir / f"doc_{cache_key}.npy"
+    
+    # Check cache first
+    if cache_file.exists():
+        try:
+            return np.load(cache_file).tolist()
+        except Exception as e:
+            print(f"Failed to load cache file {cache_file}: {e}")
+    
+    # If not in cache, call API
+    vo = voyageai.Client(api_key=VOYAGE_API_KEY)
+    max_retries = 5
+    for attempt in range(max_retries):
+        try:
+            result = vo.embed(text_chunk, model=MODEL, input_type='document')
+            embeddings = result.embeddings
+            # Save to cache
+            np.save(cache_file, embeddings)
+            return embeddings
+        except Exception as e:
+            if attempt == max_retries - 1:
+                print(f"Failed after {max_retries} attempts: {str(e)}")
+                raise e
+            time.sleep(5)
+
+
 class VoyageEmbedder(AbsEmbedder):
     def __init__(
         self,
@@ -106,64 +175,49 @@ def encode_single_device(
     ):
         return self.encode_queries(sentences, batch_size=batch_size)
 
-    def encode_queries(self, queries: List[str], batch_size: int = 256, **kwargs) -> np.ndarray:
+
+    def encode_queries(self, queries: List[str], batch_size: int = 128, num_parallel=32, **kwargs) -> np.ndarray:
+        print('Encoding queries')
+        # Prepare chunks outside of multiprocessing
+        chunks = list(split_list(queries, batch_size))
+        
+        # Setup the pool and process chunks
+        with mp.Pool(num_parallel) as pool:
+            results = list(tqdm(
+                pool.imap(call_api_query, chunks),
+                desc="Encoding queries",
+                total=len(chunks)
+            ))
 
-        #queries = [cutoff_long_text_for_embedding_generation(query, self.encoding, cutoff=4096) for query in queries]
+        # Flatten results
         total_encoded_queries = []
-        #for query_chunks in tqdm(split_list(queries, self.encoder_batch_size), total=len(queries)//self.encoder_batch_size):
-        for query_chunks in tqdm(split_list(queries, batch_size), total=len(queries)//batch_size):
-            try:
-                encoded_queries = self.vo.embed(query_chunks, model=self.embedding_model, input_type='query')
-                encoded_queries = encoded_queries.embeddings
-            except Exception as e:
-                raise e
-                time.sleep(5)
-                encoded_queries = self.vo.embed(query_chunks, model=self.embedding_model, input_type='query')
-                encoded_queries = encoded_queries.embeddings
-
-            #encoded_queries = [query_encoding for query_encoding in encoded_queries]
-            total_encoded_queries += encoded_queries
+        for result in results:
+            total_encoded_queries.extend(result)
+        
         return np.array(total_encoded_queries)
 
-    # Write your own encoding corpus function (Returns: Document embeddings as numpy array)
-    def encode_corpus(self, corpus: List[Dict[str, str]], batch_size: int = 256, **kwargs) -> np.ndarray:
+    def encode_corpus(self, corpus: List[Dict[str, str]], batch_size: int = 128, num_parallel=32, **kwargs) -> np.ndarray:
+        print('Encoding corpus')
+        # Prepare passages outside of multiprocessing
         if isinstance(corpus[0], dict):
             passages = ['{} {}'.format(doc.get('title', ''), doc['text']).strip() for doc in corpus]
         else:
             passages = corpus
+        
+        passages = [passage[:8192*8] for passage in passages]
+        chunks = list(split_list(passages, batch_size))
+
+        with mp.Pool(num_parallel) as pool:
+            results = list(tqdm(
+                pool.imap(call_api_document, chunks),
+                desc="Encoding documents",
+                total=len(chunks)
+            ))
+        
+        # Flatten results
+        total_encoded_queries = []
+        for result in results:
+            total_encoded_queries.extend(result)
 
-        passages = [
-            passage[:8192*8] #modify for context length
-            for passage in passages
-        ]
 
-        total_encoded_passages = []
-        #for passage_chunks in tqdm(split_list(passages, self.encoder_batch_size), total=len(passages)//self.encoder_batch_size):
-        for passage_chunks in tqdm(split_list(passages, batch_size), total=len(passages)//batch_size):
-            # Create a hash of the passage chunks for the cache filename
-            chunks_str = json.dumps(passage_chunks, sort_keys=True)
-            cache_key = hashlib.md5(chunks_str.encode()).hexdigest()
-            cache_file = self.cache_dir / f"{cache_key}.npy"
-
-            if cache_file.exists():
-                # Load cached embeddings if they exist
-                self.logger.info(f"Cache hit for key: {cache_key[:8]}...")
-                encoded_passages = np.load(cache_file)
-                encoded_passages = encoded_passages.tolist()
-            else:
-                attempts = 0
-                while attempts < 5 and not cache_file.exists():
-                    try:
-                        encoded_passages = self.vo.embed(passage_chunks, model=self.embedding_model, input_type='document')
-                        encoded_passages = encoded_passages.embeddings
-                        # Cache the results
-                        np.save(cache_file, encoded_passages)
-                    except Exception as e:
-                        attempts += 1
-                        self.logger.warning(f"API call failed for key: {cache_key[:8]}... Retrying after 30s")
-                        time.sleep(5)
-                if not cache_file.exists():
-                    raise Exception(f"Failed to retrieve embeddings after 5 attempts for key: {cache_key[:8]}")
-
-            total_encoded_passages += encoded_passages
         return np.array(total_encoded_passages)