Repository: Marker-Inc-Korea/AutoRAG

Issue:
[BUG] Upstage Embedding not working
At vectordb.py, llamaindex UpstageEmbedding class is instance of OpenAIEmbedding.
So it tires to truncate token to 8000, but tiktoken cannot tokenize it because solar embedding model is not capable of tokenize in tiktoken. 

Add model_name condition at truncate openai token also.



+ 
Just change colab demo to BAAI

Discussion:

Patch:
diff --git a/autorag/utils/util.py b/autorag/utils/util.py
index 5fa508f..c3b86d4 100644
--- a/autorag/utils/util.py
+++ b/autorag/utils/util.py
@@ -334,7 +334,11 @@ def save_parquet_safe(df: pd.DataFrame, filepath: str, upsert: bool = False):
 def openai_truncate_by_token(
 	texts: List[str], token_limit: int, model_name: str
 ) -> List[str]:
-	tokenizer = tiktoken.encoding_for_model(model_name)
+	try:
+		tokenizer = tiktoken.encoding_for_model(model_name)
+	except KeyError:
+		# This is not a real OpenAI model
+		return texts
 
 	def truncate_text(text: str, limit: int, tokenizer):
 		tokens = tokenizer.encode(text)


Test patch:
diff --git a/tests/autorag/utils/test_util.py b/tests/autorag/utils/test_util.py
index cd63e1f..d8c52f6 100644
--- a/tests/autorag/utils/test_util.py
+++ b/tests/autorag/utils/test_util.py
@@ -384,6 +384,11 @@ def test_openai_truncate_by_token():
 	)
 	assert len(truncated[2]) == len(t3)
 
+	truncated = openai_truncate_by_token([t1, t3], 8000, "solar-1-mini-embedding")
+	assert len(truncated) == 2
+	assert truncated[0] == base_text * 5
+	assert truncated[1] == base_text * 20
+
 
 def test_split_dataframe():
 	df = pd.DataFrame({"a": list(range(10)), "b": list(range(10, 20))})