From ecf842e17b8465848d2d3da5cf89fbad89894611 Mon Sep 17 00:00:00 2001 From: dnth Date: Fri, 8 Nov 2024 17:03:39 +0800 Subject: [PATCH] update ollama --- xinfer/ollama/llava_phi3.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/xinfer/ollama/llava_phi3.py b/xinfer/ollama/llava_phi3.py index e2c27d0..192e7cb 100644 --- a/xinfer/ollama/llava_phi3.py +++ b/xinfer/ollama/llava_phi3.py @@ -2,7 +2,7 @@ from ..model_registry import register_model from ..models import BaseModel -from ..types import ModelInputOutput +from ..types import ModelInputOutput, Result @register_model("ollama/llava-phi3", "ollama", ModelInputOutput.IMAGE_TEXT_TO_TEXT) @@ -22,15 +22,18 @@ def load_model(self): def infer_batch(self, image: str, prompt: str): raise NotImplementedError("Ollama models do not support batch inference") - def infer(self, image: str, prompt: str) -> str: + def infer(self, image: str, text: str) -> Result: res = ollama.chat( model="llava-phi3", messages=[ { "role": "user", - "content": prompt, + "content": text, "images": [image], } ], ) - return res["message"]["content"] + + result = res["message"]["content"].strip() + + return Result(text=result)