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

refactor retieval_test, add SQl retrieval methods (#61)

Browse Source
This commit is contained in:
KevinHuSh
2024-02-08 17:01:01 +08:00
committed by GitHub
parent 0a903c7714
commit 5e0a689c43
16 changed files with 238 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
rag/llm/cv_model.py
View File

@ -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