refactor stepfun cv model (#1751)
Browse files### What problem does this PR solve?
refactor stepfun cv model
### Type of change
- [x] Refactoring
Co-authored-by: Zhedong Cen <cenzhedong2@126.com>
- rag/llm/cv_model.py +1 -15
rag/llm/cv_model.py
CHANGED
|
@@ -622,27 +622,13 @@ class NvidiaCV(Base):
|
|
| 622 |
}
|
| 623 |
]
|
| 624 |
|
| 625 |
-
class StepFunCV(
|
| 626 |
def __init__(self, key, model_name="step-1v-8k", lang="Chinese", base_url="https://api.stepfun.com/v1"):
|
| 627 |
if not base_url: base_url="https://api.stepfun.com/v1"
|
| 628 |
self.client = OpenAI(api_key=key, base_url=base_url)
|
| 629 |
self.model_name = model_name
|
| 630 |
self.lang = lang
|
| 631 |
|
| 632 |
-
def describe(self, image, max_tokens=4096):
|
| 633 |
-
b64 = self.image2base64(image)
|
| 634 |
-
prompt = self.prompt(b64)
|
| 635 |
-
for i in range(len(prompt)):
|
| 636 |
-
for c in prompt[i]["content"]:
|
| 637 |
-
if "text" in c: c["type"] = "text"
|
| 638 |
-
|
| 639 |
-
res = self.client.chat.completions.create(
|
| 640 |
-
model=self.model_name,
|
| 641 |
-
messages=prompt,
|
| 642 |
-
max_tokens=max_tokens,
|
| 643 |
-
)
|
| 644 |
-
return res.choices[0].message.content.strip(), res.usage.total_tokens
|
| 645 |
-
|
| 646 |
class LmStudioCV(GptV4):
|
| 647 |
def __init__(self, key, model_name, base_url, lang="Chinese"):
|
| 648 |
if not base_url:
|
|
|
|
| 622 |
}
|
| 623 |
]
|
| 624 |
|
| 625 |
+
class StepFunCV(GptV4):
|
| 626 |
def __init__(self, key, model_name="step-1v-8k", lang="Chinese", base_url="https://api.stepfun.com/v1"):
|
| 627 |
if not base_url: base_url="https://api.stepfun.com/v1"
|
| 628 |
self.client = OpenAI(api_key=key, base_url=base_url)
|
| 629 |
self.model_name = model_name
|
| 630 |
self.lang = lang
|
| 631 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 632 |
class LmStudioCV(GptV4):
|
| 633 |
def __init__(self, key, model_name, base_url, lang="Chinese"):
|
| 634 |
if not base_url:
|