fix: Implement correct batch handling for Ollama embedding and improve progress updates

rag/flow/tokenizer/tokenizer.py

@@ -15,6 +15,7 @@
 import logging
 import random
 import re
+import time
 
 import numpy as np
 
@@ -84,16 +85,21 @@ def batch_encode(txts):
             return embedding_model.encode([truncate(c, embedding_model.max_length - 10) for c in txts])
 
         cnts_ = np.array([])
+        self.callback(1, f"Total chunks to embed: {len(texts)}")
+
         for i in range(0, len(texts), settings.EMBEDDING_BATCH_SIZE):
             async with embed_limiter:
+                start = time.time()
                 vts, c = await thread_pool_exec(batch_encode,texts[i : i + settings.EMBEDDING_BATCH_SIZE],)
+                delta = time.time()-start
             if len(cnts_) == 0:
                 cnts_ = vts
             else:
                 cnts_ = np.concatenate((cnts_, vts), axis=0)
             token_count += c
             if i % 33 == 32:
                 self.callback(i * 1.0 / len(texts) / parts / settings.EMBEDDING_BATCH_SIZE + 0.5 * (parts - 1))
+            self.callback(1, f"Embedding batch {i}-{i + settings.EMBEDDING_BATCH_SIZE} :: {(i*100)/len(texts):2.2f}% done :: {settings.EMBEDDING_BATCH_SIZE / delta:2.1f} chunks/s")
 
         cnts = cnts_
         title_w = float(self._param.filename_embd_weight)

rag/llm/embedding_model.py

@@ -256,6 +256,9 @@ def encode_queries(self, text):
             raise Exception(f"Error: {res}")
 
 
+
+
+
 class OllamaEmbed(Base):
     _FACTORY_NAME = "Ollama"
 
@@ -267,27 +270,32 @@ def __init__(self, key, model_name, **kwargs):
         self.keep_alive = kwargs.get("ollama_keep_alive", int(os.environ.get("OLLAMA_KEEP_ALIVE", -1)))
 
     def encode(self, texts: list):
-        arr = []
+        batch = []
         tks_num = 0
+        # Remove special tokens if they exist
         for txt in texts:
             # remove special tokens if they exist base on regex in one request
+            mod_txt = txt
             for token in OllamaEmbed._special_tokens:
-                txt = txt.replace(token, "")
-            res = self.client.embeddings(prompt=txt, model=self.model_name, options={"use_mmap": True}, keep_alive=self.keep_alive)
-            try:
-                arr.append(res["embedding"])
-            except Exception as _e:
-                log_exception(_e, res)
-                raise Exception(f"Error: {res}")
-            tks_num += 128
-        return np.array(arr), tks_num
+                mod_txt = mod_txt.replace(token, "")
+            batch.append(mod_txt)
+        embeddings = []
+
+        try:
+            res = self.client.embed(input=batch, model=self.model_name, options={"use_mmap": True}, keep_alive=self.keep_alive)
+            embeddings.extend(res["embeddings"])
+        except Exception as _e:
+            log_exception(_e, res)
+            raise Exception(f"Error: {_e}")
+        tks_num += 128
+        return np.array(embeddings), 128
 
     def encode_queries(self, text):
         # remove special tokens if they exist
         for token in OllamaEmbed._special_tokens:
             text = text.replace(token, "")
-        res = self.client.embeddings(prompt=text, model=self.model_name, options={"use_mmap": True}, keep_alive=self.keep_alive)
         try:
+            res = self.client.embeddings(prompt=text, model=self.model_name, options={"use_mmap": True}, keep_alive=self.keep_alive)
             return np.array(res["embedding"]), 128
         except Exception as _e:
             log_exception(_e, res)