|
|
from abc import ABC |
|
|
from openai import OpenAI |
|
|
import os |
|
|
import base64 |
|
|
from io import BytesIO |
|
|
|
|
|
class Base(ABC): |
|
|
def describe(self, image, max_tokens=300): |
|
|
raise NotImplementedError("Please implement encode method!") |
|
|
|
|
|
|
|
|
class GptV4(Base): |
|
|
def __init__(self): |
|
|
import openapi |
|
|
openapi.api_key = os.environ["OPENAPI_KEY"] |
|
|
self.client = OpenAI() |
|
|
|
|
|
def describe(self, image, max_tokens=300): |
|
|
buffered = BytesIO() |
|
|
try: |
|
|
image.save(buffered, format="JPEG") |
|
|
except Exception as e: |
|
|
image.save(buffered, format="PNG") |
|
|
b64 = base64.b64encode(buffered.getvalue()).decode("utf-8") |
|
|
|
|
|
res = self.client.chat.completions.create( |
|
|
model="gpt-4-vision-preview", |
|
|
messages=[ |
|
|
{ |
|
|
"role": "user", |
|
|
"content": [ |
|
|
{ |
|
|
"type": "text", |
|
|
"text": "请用中文详细描述一下图中的内容，比如时间，地点，人物，事情，人物心情等。", |
|
|
}, |
|
|
{ |
|
|
"type": "image_url", |
|
|
"image_url": { |
|
|
"url": f"data:image/jpeg;base64,{b64}" |
|
|
}, |
|
|
}, |
|
|
], |
|
|
} |
|
|
], |
|
|
max_tokens=max_tokens, |
|
|
) |
|
|
return res.choices[0].message.content.strip() |
|
|
|
|
|
|
|
|
class QWen(Base): |
|
|
def chat(self, system, history, gen_conf): |
|
|
from http import HTTPStatus |
|
|
from dashscope import Generation |
|
|
from dashscope.api_entities.dashscope_response import Role |
|
|
|
|
|
response = Generation.call( |
|
|
Generation.Models.qwen_turbo, |
|
|
messages=messages, |
|
|
result_format='message' |
|
|
) |
|
|
if response.status_code == HTTPStatus.OK: |
|
|
return response.output.choices[0]['message']['content'] |
|
|
return response.message |
|
|
|
