Update rag/RAG.py

rag/RAG.py

@@ -27,32 +27,40 @@ class RAG:
         except Exception as e: 
             raise ValueError(f"an error occured: {e}")
     
+    def _embed_batch(self, chunk_batch, task_type):
+        response = client.models.embed_content(
+            model=self.MODEL,
+            contents=chunk_batch,
+            config=types.EmbedContentConfig(task_type=task_type)
+        )
+        return [embedding.values for embedding in response.embeddings]
+    
     def generate_embedding(self, text, task_type=None):
         try:
             if not task_type:
                 task_type = self.TASK_TYPE
-    
+
             chunks = self.split_text(text)
-            batches = [chunks[i:i + self.MAX_BATCH_SIZE] for i in range(0, len(chunks), self.MAX_BATCH_SIZE)]
-    
-            def embed_batch(batch):
-                response = client.models.embed_content(
-                    model=self.MODEL,
-                    contents=batch,
-                    config=types.EmbedContentConfig(task_type=task_type)
-                )
-                return [embedding.values for embedding in response.embeddings]
-    
+            batches = [
+                chunks[i:i + self.MAX_BATCH_SIZE]
+                for i in range(0, len(chunks), self.MAX_BATCH_SIZE)
+            ]
+
             embeddings = []
-            with ThreadPoolExecutor(max_workers=100) as executor:
-                futures = [executor.submit(embed_batch, batch) for batch in batches]
+            with ThreadPoolExecutor(max_workers=50) as executor:
+                futures = {
+                    executor.submit(self._embed_batch, batch, task_type): batch
+                    for batch in batches
+                }
+
                 for future in as_completed(futures):
-                    embeddings.extend(future.result())
-    
+                    result = future.result()
+                    embeddings.extend(result)
+
             return {"embeddings": embeddings, "chunks": chunks}, 200
-    
+
         except Exception as e:
-            return {"an error occurred": str(e)}, 500
+            return {"an error occurred": f"{e}"}, 500