AI/ragflow
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/infiniflow/ragflow.git synced 2025-08-14 16:55:53 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor stepfun cv model (#1751)

Browse Source 
### What problem does this PR solve?

refactor stepfun cv model

### Type of change

- [x] Refactoring

Co-authored-by: Zhedong Cen <cenzhedong2@126.com>
This commit is contained in:
黄腾 2024-07-31 15:30:47 +08:00 committed by GitHub
parent 9169643157
commit ccbd4365be
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
rag/llm/cv_model.py
View File

@ -622,27 +622,13 @@ class NvidiaCV(Base):
} }
] ]
class StepFunCV(Base): class StepFunCV(GptV4):
def __init__(self, key, model_name="step-1v-8k", lang="Chinese", base_url="https://api.stepfun.com/v1"): def __init__(self, key, model_name="step-1v-8k", lang="Chinese", base_url="https://api.stepfun.com/v1"):
if not base_url: base_url="https://api.stepfun.com/v1" if not base_url: base_url="https://api.stepfun.com/v1"
self.client = OpenAI(api_key=key, base_url=base_url) self.client = OpenAI(api_key=key, base_url=base_url)
self.model_name = model_name self.model_name = model_name
self.lang = lang self.lang = lang
def describe(self, image, max_tokens=4096):
b64 = self.image2base64(image)
prompt = self.prompt(b64)
for i in range(len(prompt)):
for c in prompt[i]["content"]:
if "text" in c: c["type"] = "text"
res = self.client.chat.completions.create(
model=self.model_name,
messages=prompt,
max_tokens=max_tokens,
)
return res.choices[0].message.content.strip(), res.usage.total_tokens
class LmStudioCV(GptV4): class LmStudioCV(GptV4):
def __init__(self, key, model_name, base_url, lang="Chinese"): def __init__(self, key, model_name, base_url, lang="Chinese"):
if not base_url: if not base_url: