File size: 1,023 Bytes
58974f8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
from models import BaseModel
from .base_refiner import BaseRefiner
from utils.image_encoder import encode_image
class SimpleRefiner(BaseRefiner):
def __init__(self,
sys_prompt: str,
model: BaseModel,
) -> None:
BaseRefiner.__init__(self, sys_prompt=sys_prompt, model=model)
def refine(self, message: str, image_path=None) -> str:
if image_path:
context = [{"role": "system", "content": self.sys_prompt}, {"role": "user", "content": [
{"type": "text", "text": f"{message}"},
{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encode_image(image_path)}"}},
]}]
else:
context = [{"role": "system", "content": self.sys_prompt}, {"role": "user", "content": [
{"type": "text", "text": f"{message}"},
]}]
response = self.model.respond(context)
return response |