refactor retieval_test, add SQl retrieval methods (#61)

rag/llm/chat_model.py

@@ -58,3 +58,21 @@ class QWenChat(Base):
         if response.status_code == HTTPStatus.OK:
             return response.output.choices[0]['message']['content'], response.usage.output_tokens
         return response.message, 0
+
+
+from zhipuai import ZhipuAI
+class ZhipuChat(Base):
+    def __init__(self, key, model_name="glm-3-turbo"):
+        self.client = ZhipuAI(api_key=key)
+        self.model_name = model_name
+
+    def chat(self, system, history, gen_conf):
+        from http import HTTPStatus
+        history.insert(0, {"role": "system", "content": system})
+        response = self.client.chat.completions.create(
+            self.model_name,
+            messages=history
+        )
+        if response.status_code == HTTPStatus.OK:
+            return response.output.choices[0]['message']['content'], response.usage.completion_tokens
+        return response.message, 0

rag/llm/cv_model.py

@@ -61,7 +61,7 @@ class Base(ABC):
 
 class GptV4(Base):
     def __init__(self, key, model_name="gpt-4-vision-preview"):
-        self.client = OpenAI(api_key = key)
+        self.client = OpenAI(api_key=key)
         self.model_name = model_name
 
     def describe(self, image, max_tokens=300):
@@ -89,3 +89,22 @@ class QWenCV(Base):
         if response.status_code == HTTPStatus.OK:
             return response.output.choices[0]['message']['content'], response.usage.output_tokens
         return response.message, 0
+
+
+from zhipuai import ZhipuAI
+
+
+class Zhipu4V(Base):
+    def __init__(self, key, model_name="glm-4v"):
+        self.client = ZhipuAI(api_key=key)
+        self.model_name = model_name
+
+    def describe(self, image, max_tokens=1024):
+        b64 = self.image2base64(image)
+
+        res = self.client.chat.completions.create(
+            model=self.model_name,
+            messages=self.prompt(b64),
+            max_tokens=max_tokens,
+        )
+        return res.choices[0].message.content.strip(), res.usage.total_tokens

rag/llm/embedding_model.py

@@ -19,7 +19,6 @@ import dashscope
 from openai import OpenAI
 from FlagEmbedding import FlagModel
 import torch
-import os
 import numpy as np
 
 from rag.utils import num_tokens_from_string
@@ -114,4 +113,21 @@ class QWenEmbed(Base):
                 input=text[:2048],
                 text_type="query"
             )
-        return np.array(resp["output"]["embeddings"][0]["embedding"]), resp["usage"]["input_tokens"]
+        return np.array(resp["output"]["embeddings"][0]["embedding"]), resp["usage"]["input_tokens"]
+
+
+from zhipuai import ZhipuAI
+class ZhipuEmbed(Base):
+    def __init__(self, key, model_name="embedding-2"):
+        self.client = ZhipuAI(api_key=key)
+        self.model_name = model_name
+
+    def encode(self, texts: list, batch_size=32):
+        res = self.client.embeddings.create(input=texts,
+                                            model=self.model_name)
+        return np.array([d.embedding for d in res.data]), res.usage.total_tokens
+
+    def encode_queries(self, text):
+        res = self.client.embeddings.create(input=text,
+                                            model=self.model_name)
+        return np.array(res["data"][0]["embedding"]), res.usage.total_tokens