黄腾 aopstudio commited on
Commit
c579820
·
1 Parent(s): 479edea

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]>

Files changed (1) hide show
  1. 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(Base):
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: