update ollama

dnth · Nov 8, 2024 · ecf842e · ecf842e
1 parent 72daf0e
commit ecf842e
Showing 1 changed file with 7 additions and 4 deletions.
diff --git 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)