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 <[email protected]>
- 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:
|