diff --git a/rag/llm/__init__.py b/rag/llm/__init__.py index ce78e2914..323ac9502 100644 --- a/rag/llm/__init__.py +++ b/rag/llm/__init__.py @@ -45,6 +45,7 @@ from .embedding_model import ( HuggingFaceEmbed, VolcEngineEmbed, GPUStackEmbed, + NovitaEmbed ) from .chat_model import ( GptTurbo, @@ -128,6 +129,7 @@ from .rerank_model import ( QWenRerank, GPUStackRerank, HuggingfaceRerank, + NovitaRerank ) from .sequence2txt_model import ( @@ -180,6 +182,7 @@ EmbeddingModel = { "HuggingFace": HuggingFaceEmbed, "VolcEngine": VolcEngineEmbed, "GPUStack": GPUStackEmbed, + "NovitaAI": NovitaEmbed } CvModel = { @@ -267,6 +270,7 @@ RerankModel = { "Tongyi-Qianwen": QWenRerank, "GPUStack": GPUStackRerank, "HuggingFace": HuggingfaceRerank, + "NovitaAI": NovitaRerank } Seq2txtModel = { diff --git a/rag/llm/embedding_model.py b/rag/llm/embedding_model.py index 4b1a48abe..b9a40bbd0 100644 --- a/rag/llm/embedding_model.py +++ b/rag/llm/embedding_model.py @@ -913,3 +913,8 @@ class GPUStackEmbed(OpenAIEmbed): self.client = OpenAI(api_key=key, base_url=base_url) self.model_name = model_name + + +class NovitaEmbed(SILICONFLOWEmbed): + def __init__(self, key, model_name, base_url="https://api.novita.ai/v3/openai/embeddings"): + super().__init__(key, model_name, base_url) \ No newline at end of file diff --git a/rag/llm/rerank_model.py b/rag/llm/rerank_model.py index 45ef4bb0f..5310a5f3b 100644 --- a/rag/llm/rerank_model.py +++ b/rag/llm/rerank_model.py @@ -626,3 +626,7 @@ class GPUStackRerank(Base): raise ValueError( f"Error calling GPUStackRerank model {self.model_name}: {e.response.status_code} - {e.response.text}") + +class NovitaRerank(JinaRerank): + def __init__(self, key, model_name, base_url="https://api.novita.ai/v3/openai/rerank"): + super().__init__(key, model_name, base_url) \ No newline at end of file