add support for NVIDIA llm (#1645)

### What problem does this PR solve?

add support for NVIDIA llm
### Type of change

- [x] New Feature (non-breaking change which adds functionality)

---------

Co-authored-by: Zhedong Cen <cenzhedong2@126.com>

rag/llm/__init__.py

@@ -34,7 +34,8 @@ EmbeddingModel = {
     "BAAI": DefaultEmbedding,
     "Mistral": MistralEmbed,
     "Bedrock": BedrockEmbed,
-    "Gemini":GeminiEmbed
+    "Gemini":GeminiEmbed,
+    "NVIDIA":NvidiaEmbed
 }
 
 
@@ -48,7 +49,8 @@ CvModel = {
     "Moonshot": LocalCV,
     'Gemini':GeminiCV,
     'OpenRouter':OpenRouterCV,
-    "LocalAI":LocalAICV
+    "LocalAI":LocalAICV,
+    "NVIDIA":NvidiaCV
 }
 
 
@@ -71,7 +73,8 @@ ChatModel = {
     "Bedrock": BedrockChat,
     "Groq": GroqChat,
     'OpenRouter':OpenRouterChat,
-    "StepFun":StepFunChat
+    "StepFun":StepFunChat,
+    "NVIDIA":NvidiaChat
 }
 
 
@@ -79,7 +82,8 @@ RerankModel = {
     "BAAI": DefaultRerank,
     "Jina": JinaRerank,
     "Youdao": YoudaoRerank,
-    "Xinference": XInferenceRerank
+    "Xinference": XInferenceRerank,
+    "NVIDIA":NvidiaRerank
 }
 
 

rag/llm/chat_model.py

@@ -581,7 +581,6 @@ class MiniMaxChat(Base):
             response = requests.request(
                 "POST", url=self.base_url, headers=headers, data=payload
             )
-            print(response, flush=True)
             response = response.json()
             ans = response["choices"][0]["message"]["content"].strip()
             if response["choices"][0]["finish_reason"] == "length":
@@ -902,4 +901,79 @@ class StepFunChat(Base):
     def __init__(self, key, model_name, base_url="https://api.stepfun.com/v1/chat/completions"):
         if not base_url:
             base_url = "https://api.stepfun.com/v1/chat/completions"
-        super().__init__(key, model_name, base_url)
+        super().__init__(key, model_name, base_url)
+
+
+class NvidiaChat(Base):
+    def __init__(
+        self,
+        key,
+        model_name,
+        base_url="https://integrate.api.nvidia.com/v1/chat/completions",
+    ):
+        if not base_url:
+            base_url = "https://integrate.api.nvidia.com/v1/chat/completions"
+        self.base_url = base_url
+        self.model_name = model_name
+        self.api_key = key
+        self.headers = {
+            "accept": "application/json",
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+        }
+
+    def chat(self, system, history, gen_conf):
+        if system:
+            history.insert(0, {"role": "system", "content": system})
+        for k in list(gen_conf.keys()):
+            if k not in ["temperature", "top_p", "max_tokens"]:
+                del gen_conf[k]
+        payload = {"model": self.model_name, "messages": history, **gen_conf}
+        try:
+            response = requests.post(
+                url=self.base_url, headers=self.headers, json=payload
+            )
+            response = response.json()
+            ans = response["choices"][0]["message"]["content"].strip()
+            return ans, response["usage"]["total_tokens"]
+        except Exception as e:
+            return "**ERROR**: " + str(e), 0
+
+    def chat_streamly(self, system, history, gen_conf):
+        if system:
+            history.insert(0, {"role": "system", "content": system})
+        for k in list(gen_conf.keys()):
+            if k not in ["temperature", "top_p", "max_tokens"]:
+                del gen_conf[k]
+        ans = ""
+        total_tokens = 0
+        payload = {
+            "model": self.model_name,
+            "messages": history,
+            "stream": True,
+            **gen_conf,
+        }
+
+        try:
+            response = requests.post(
+                url=self.base_url,
+                headers=self.headers,
+                json=payload,
+            )
+            for resp in response.text.split("\n\n"):
+                if "choices" not in resp:
+                    continue
+                resp = json.loads(resp[6:])
+                if "content" in resp["choices"][0]["delta"]:
+                    text = resp["choices"][0]["delta"]["content"]
+                else:
+                    continue
+                ans += text
+                if "usage" in resp:
+                    total_tokens = resp["usage"]["total_tokens"]
+                yield ans
+
+        except Exception as e:
+            yield ans + "\n**ERROR**: " + str(e)
+
+        yield total_tokens

rag/llm/cv_model.py

@@ -137,7 +137,6 @@ class Base(ABC):
         ]
 
 
-
 class GptV4(Base):
     def __init__(self, key, model_name="gpt-4-vision-preview", lang="Chinese", base_url="https://api.openai.com/v1"):
         if not base_url: base_url="https://api.openai.com/v1"
@@ -619,3 +618,65 @@ class LocalCV(Base):
 
     def describe(self, image, max_tokens=1024):
         return "", 0
+
+
+class NvidiaCV(Base):
+    def __init__(
+        self,
+        key,
+        model_name,
+        lang="Chinese",
+        base_url="https://ai.api.nvidia.com/v1/vlm",
+    ):
+        if not base_url:
+            base_url = ("https://ai.api.nvidia.com/v1/vlm",)
+        self.lang = lang
+        factory, llm_name = model_name.split("/")
+        if factory != "liuhaotian":
+            self.base_url = os.path.join(base_url, factory, llm_name)
+        else:
+            self.base_url = os.path.join(
+                base_url, "community", llm_name.replace("-v1.6", "16")
+            )
+        self.key = key
+
+    def describe(self, image, max_tokens=1024):
+        b64 = self.image2base64(image)
+        response = requests.post(
+            url=self.base_url,
+            headers={
+                "accept": "application/json",
+                "content-type": "application/json",
+                "Authorization": f"Bearer {self.key}",
+            },
+            json={
+                "messages": self.prompt(b64),
+                "max_tokens": max_tokens,
+            },
+        )
+        response = response.json()
+        return (
+            response["choices"][0]["message"]["content"].strip(),
+            response["usage"]["total_tokens"],
+        )
+
+    def prompt(self, b64):
+        return [
+            {
+                "role": "user",
+                "content": (
+                    "请用中文详细描述一下图中的内容，比如时间，地点，人物，事情，人物心情等，如果有数据请提取出数据。"
+                    if self.lang.lower() == "chinese"
+                    else "Please describe the content of this picture, like where, when, who, what happen. If it has number data, please extract them out."
+                )
+                + f' <img src="data:image/jpeg;base64,{b64}"/>',
+            }
+        ]
+
+    def chat_prompt(self, text, b64):
+        return [
+            {
+                "role": "user",
+                "content": text + f' <img src="data:image/jpeg;base64,{b64}"/>',
+            }
+        ]

rag/llm/embedding_model.py

@@ -462,3 +462,41 @@ class GeminiEmbed(Base):
             title="Embedding of single string")
         token_count = num_tokens_from_string(text)
         return np.array(result['embedding']),token_count
+
+class NvidiaEmbed(Base):
+    def __init__(
+        self, key, model_name, base_url="https://integrate.api.nvidia.com/v1/embeddings"
+    ):
+        if not base_url:
+            base_url = "https://integrate.api.nvidia.com/v1/embeddings"
+        self.api_key = key
+        self.base_url = base_url
+        self.headers = {
+            "accept": "application/json",
+            "Content-Type": "application/json",
+            "authorization": f"Bearer {self.api_key}",
+        }
+        self.model_name = model_name
+        if model_name == "nvidia/embed-qa-4":
+            self.base_url = "https://ai.api.nvidia.com/v1/retrieval/nvidia/embeddings"
+            self.model_name = "NV-Embed-QA"
+        if model_name == "snowflake/arctic-embed-l":
+            self.base_url = "https://ai.api.nvidia.com/v1/retrieval/snowflake/arctic-embed-l/embeddings"
+
+    def encode(self, texts: list, batch_size=None):
+        payload = {
+            "input": texts,
+            "input_type": "query",
+            "model": self.model_name,
+            "encoding_format": "float",
+            "truncate": "END",
+        }
+        res = requests.post(self.base_url, headers=self.headers, json=payload).json()
+        return (
+            np.array([d["embedding"] for d in res["data"]]),
+            res["usage"]["total_tokens"],
+        )
+
+    def encode_queries(self, text):
+        embds, cnt = self.encode([text])
+        return np.array(embds[0]), cnt

rag/llm/rerank_model.py

@@ -164,3 +164,41 @@ class LocalAIRerank(Base):
 
     def similarity(self, query: str, texts: list):
         raise NotImplementedError("The LocalAIRerank has not been implement")
+
+
+class NvidiaRerank(Base):
+    def __init__(
+        self, key, model_name, base_url="https://ai.api.nvidia.com/v1/retrieval/nvidia/"
+    ):
+        if not base_url:
+            base_url = "https://ai.api.nvidia.com/v1/retrieval/nvidia/"
+        self.model_name = model_name
+
+        if self.model_name == "nvidia/nv-rerankqa-mistral-4b-v3":
+            self.base_url = os.path.join(
+                base_url, "nv-rerankqa-mistral-4b-v3", "reranking"
+            )
+
+        if self.model_name == "nvidia/rerank-qa-mistral-4b":
+            self.base_url = os.path.join(base_url, "reranking")
+            self.model_name = "nv-rerank-qa-mistral-4b:1"
+
+        self.headers = {
+            "accept": "application/json",
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {key}",
+        }
+
+    def similarity(self, query: str, texts: list):
+        token_count = num_tokens_from_string(query) + sum(
+            [num_tokens_from_string(t) for t in texts]
+        )
+        data = {
+            "model": self.model_name,
+            "query": {"text": query},
+            "passages": [{"text": text} for text in texts],
+            "truncate": "END",
+            "top_n": len(texts),
+        }
+        res = requests.post(self.base_url, headers=self.headers, json=data).json()
+        return (np.array([d["logit"] for d in res["rankings"]]), token_count)