Ver código fonte

add step-1v-8k cv model (#1686)

### What problem does this PR solve?

_Briefly describe what this PR aims to solve. Include background context
that will help reviewers understand the purpose of the PR._

### Type of change

- [x] New Feature (non-breaking change which adds functionality)

---------

Co-authored-by: lijianyong <lijianyong@stepfun.com>
tags/v0.9.0
jianyongli 1 ano atrás
pai
commit
9169643157
Nenhuma conta vinculada ao e-mail do autor do commit
3 arquivos alterados com 23 adições e 2 exclusões
  1. 1
    1
      conf/llm_factories.json
  2. 2
    1
      rag/llm/__init__.py
  3. 20
    0
      rag/llm/cv_model.py

+ 1
- 1
conf/llm_factories.json Ver arquivo

@@ -1920,7 +1920,7 @@
{
"llm_name": "step-1v-8k",
"tags": "LLM,CHAT,IMAGE2TEXT",
"max_tokens": 8000,
"max_tokens": 8192,
"model_type": "image2text"
}
]

+ 2
- 1
rag/llm/__init__.py Ver arquivo

@@ -52,7 +52,8 @@ CvModel = {
"OpenRouter": OpenRouterCV,
"LocalAI": LocalAICV,
"NVIDIA": NvidiaCV,
"LM-Studio": LmStudioCV
"LM-Studio": LmStudioCV,
"StepFun":StepFunCV
}



+ 20
- 0
rag/llm/cv_model.py Ver arquivo

@@ -622,6 +622,26 @@ class NvidiaCV(Base):
}
]

class StepFunCV(Base):
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"):

Carregando…
Cancelar
Salvar