Handle azure models as well

topoteretes · Jan 7, 2025 · fb13a1b · fb13a1b
1 parent a774191
commit fb13a1b
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/cognee/tasks/chunks/chunk_by_paragraph.py b/cognee/tasks/chunks/chunk_by_paragraph.py
@@ -29,6 +29,8 @@ def chunk_by_paragraph(
     for paragraph_id, sentence, word_count, end_type in chunk_by_sentence(data, maximum_length=paragraph_length):
         # Check if this sentence would exceed length limit
         if embedding_model:
+            if embedding_model.startswith("azure/"):
+                embedding_model = embedding_model.split("/")[-1]
             tokenizer = tiktoken.encoding_for_model(embedding_model)
             token_count = len(tokenizer.encode(sentence))
         else: