pengjunjie/ragflow
1
0
Fork 1
mirror of https://github.com/infiniflow/ragflow.git synced 2025-12-08 20:42:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

Make spark model robuster to model name (#3514)

Browse Source 
### What problem does this PR solve?


### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
This commit is contained in:
Kevin Hu
2024-11-20 12:16:36 +08:00
committed by Yingfeng Zhang
parent cad341e794
commit 81c7b6afc5
2 changed files with 9 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
rag/llm/chat_model.py
View File

@ -1164,7 +1164,11 @@ class SparkChat(Base):
"Spark-Pro-128K": "pro-128k",
"Spark-4.0-Ultra": "4.0Ultra",
}
model_version = model2version[model_name]
version2model = {v: k for k, v in model2version.items()}
assert model_name in model2version or model_name in version2model, f"The given model name is not supported yet. Support: {list(model2version.keys())}"
if model_name in model2version:
model_version = model2version[model_name]
else: model_version = model_name
super().__init__(key, model_version, base_url)