Merge pull request #3574 from dengpeng777/master

support zhipu-api embedding

server/model_workers/zhipu.py

@@ -39,6 +39,8 @@ def generate_token(apikey: str, exp_seconds: int):
 
 
 class ChatGLMWorker(ApiModelWorker):
+ DEFAULT_EMBED_MODEL = "embedding-2"
+
  def __init__(
  self,
  *,
@@ -84,7 +86,29 @@ def do_chat(self, params: ApiChatParams) -> Iterator[Dict]:
  # yield {"error_code": 0, "text": text}
 
 
-
+ def do_embeddings(self, params: ApiEmbeddingsParams) -> Dict:
+ params.load_config(self.model_names[0])
+ token = generate_token(params.api_key, 60)
+ headers = {
+ "Content-Type": "application/json",
+ "Authorization": f"Bearer {token}"
+ }
+ i = 0
+ batch_size = 1
+ result = []
+ while i < len(params.texts):
+ data = {
+ "model": params.embed_model or self.DEFAULT_EMBED_MODEL,
+ "input": "".join(params.texts[i: i + batch_size])
+ }
+ url = "https://open.bigmodel.cn/api/paas/v4/embeddings"
+ response = requests.post(url, headers=headers, json=data)
+ ans = response.json()
+ result.append(ans["data"][0]["embedding"])
+ i += batch_size
+
+ return {"code": 200, "data": result}
+
  def get_embeddings(self, params):
  print("embedding")
  print(params)