Spaces:
Sleeping
Sleeping
from models import BaseModel | |
from .base_refiner import BaseRefiner | |
from utils.image_encoder import encode_image | |
import asyncio | |
class SimpleRefiner(BaseRefiner): | |
def __init__(self, | |
sys_prompt: str, | |
model: BaseModel, | |
) -> None: | |
BaseRefiner.__init__(self, sys_prompt=sys_prompt, model=model) | |
async def refine_async(self, message: str, memory, image_paths=None) -> str: | |
if memory is None: | |
memory = [] | |
else: | |
memory = memory.messages[1:] | |
user_context = [{"role": "user", "content": [ | |
{"type": "text", "text": f"{message}"},]}] | |
if image_paths: | |
if not isinstance(image_paths, list): | |
image_paths = [image_paths] | |
for image_path in image_paths: | |
user_context[0]["content"].append({"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encode_image(image_path.name)}"}}) | |
context = [{"role": "system", "content": self.sys_prompt}] + memory + user_context | |
respond_task = asyncio.create_task(self.model.respond_async(context)) | |
await respond_task | |
response = respond_task.result() | |
return response | |
def refine(self, message: str, memory, image_paths=None) -> str: | |
if memory is None: | |
memory = [] | |
else: | |
memory = memory.messages[1:] | |
user_context = [{"role": "user", "content": [ | |
{"type": "text", "text": f"{message}"},]}] | |
if image_paths: | |
if not isinstance(image_paths, list): | |
image_paths = [image_paths] | |
for image_path in image_paths: | |
user_context[0]["content"].append({"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encode_image(image_path.name)}"}}) | |
context = [{"role": "system", "content": self.sys_prompt}] + memory + user_context | |
response = self.model.respond(context) | |
return response |