Support Ollama (#261)

### What problem does this PR solve?

Issue link:#221

### Type of change

- [x] New Feature (non-breaking change which adds functionality)
This commit is contained in:
KevinHuSh
2024-04-08 19:20:57 +08:00
committed by GitHub
parent 265a7a283a
commit 3708b97db9
15 changed files with 234 additions and 43 deletions

View File

@ -16,7 +16,7 @@
from zhipuai import ZhipuAI
import io
from abc import ABC
from ollama import Client
from PIL import Image
from openai import OpenAI
import os
@ -140,6 +140,28 @@ class Zhipu4V(Base):
return res.choices[0].message.content.strip(), res.usage.total_tokens
class OllamaCV(Base):
def __init__(self, key, model_name, lang="Chinese", **kwargs):
self.client = Client(host=kwargs["base_url"])
self.model_name = model_name
self.lang = lang
def describe(self, image, max_tokens=1024):
prompt = self.prompt("")
try:
options = {"num_predict": max_tokens}
response = self.client.generate(
model=self.model_name,
prompt=prompt[0]["content"][1]["text"],
images=[image],
options=options
)
ans = response["response"].strip()
return ans, 128
except Exception as e:
return "**ERROR**: " + str(e), 0
class LocalCV(Base):
def __init__(self, key, model_name="glm-4v", lang="Chinese", **kwargs):
pass