From ccbd4365be17c874d5721841ee952203c1357ba4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=BB=84=E8=85=BE?= <101850389+hangters@users.noreply.github.com> Date: Wed, 31 Jul 2024 15:30:47 +0800 Subject: [PATCH] refactor stepfun cv model (#1751) ### What problem does this PR solve? refactor stepfun cv model ### Type of change - [x] Refactoring Co-authored-by: Zhedong Cen --- rag/llm/cv_model.py | 16 +--------------- 1 file changed, 1 insertion(+), 15 deletions(-) diff --git a/rag/llm/cv_model.py b/rag/llm/cv_model.py index baf93fd89..1277791ee 100644 --- a/rag/llm/cv_model.py +++ b/rag/llm/cv_model.py @@ -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"): if not base_url: base_url="https://api.stepfun.com/v1" self.client = OpenAI(api_key=key, base_url=base_url) self.model_name = model_name 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): def __init__(self, key, model_name, base_url, lang="Chinese"): if not base_url: