Fix Ollama inference engine

.github/workflows/inference_tests.yml

@@ -11,40 +11,46 @@ on:
     types: [published]
 
 concurrency:
-    group: ${{ github.workflow }}-${{ github.event_name == 'pull_request' && github.event.pull_request.number || github.ref_name }}
-    cancel-in-progress: true
+  group: ${{ github.workflow }}-${{ github.event_name == 'pull_request' && github.event.pull_request.number || github.ref_name }}
+  cancel-in-progress: true
 
 jobs:
-   inference:
+  inference:
+    runs-on: ubuntu-latest
+    env:
+      OS: ubuntu-latest
+      UNITXT_DEFAULT_VERBOSITY: error
+      DATASETS_VERBOSITY: error
+      HF_HUB_VERBOSITY: error
+      HF_DATASETS_DISABLE_PROGRESS_BARS: "True"
+      TQDM_DISABLE: "True"
+      WML_URL: ${{ secrets.WML_URL }}
+      WML_PROJECT_ID: ${{ secrets.WML_PROJECT_ID }}
+      WML_APIKEY: ${{ secrets.WML_APIKEY }}
+      WX_URL: ${{ secrets.WX_URL }}
+      WX_PROJECT_ID: ${{ secrets.WX_PROJECT_ID }}
+      WX_API_KEY: ${{ secrets.WX_API_KEY }}
+      GENAI_KEY: ${{ secrets.GENAI_KEY }}
+    steps:
+      - uses: actions/checkout@v4
 
-     runs-on: ubuntu-latest
-     env:
-       OS: ubuntu-latest
-       UNITXT_DEFAULT_VERBOSITY: error
-       DATASETS_VERBOSITY: error
-       HF_HUB_VERBOSITY: error
-       HF_DATASETS_DISABLE_PROGRESS_BARS: "True"
-       TQDM_DISABLE: "True"
-       WML_URL: ${{ secrets.WML_URL }}
-       WML_PROJECT_ID: ${{ secrets.WML_PROJECT_ID }}
-       WML_APIKEY: ${{ secrets.WML_APIKEY }}
-       WX_URL: ${{ secrets.WML_URL }}
-       WX_PROJECT_ID: ${{ secrets.WML_PROJECT_ID }}
-       WX_API_KEY: ${{ secrets.WML_APIKEY }}
-       GENAI_KEY: ${{ secrets.GENAI_KEY }}
-     steps:
-     - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.10'
 
-     - uses: actions/setup-python@v5
-       with:
-         python-version: '3.10'
-
-     - run: curl -LsSf https://astral.sh/uv/install.sh | sh
-     - run: uv pip install --system ".[tests,watsonx,inference-tests]"
-     - name:  Hugging Face Login
+      - run: curl -LsSf https://astral.sh/uv/install.sh | sh
+      - run: uv pip install --system ".[tests,watsonx,inference-tests]"
+      - name: Hugging Face Login
         run: |
           for i in {1..5}; do
             huggingface-cli login --token ${{ secrets.UNITXT_READ_HUGGINGFACE_HUB_FOR_TESTS }} && break || sleep $((2 ** i));
           done
-     - name: Run Tests
-       run: python -m unittest discover -s tests/inference -p "test_*.py"
+      - name: Install Ollama
+        run: curl -fsSL https://ollama.com/install.sh | sh
+      - name: Pull Llama 3.2:1b model
+        run: ollama pull llama3.2:1b
+      - name: Start serving the model
+        run: nohup ollama serve --model llama3.2:1b --port 5000 &
+
+      - name: Run Tests
+        run: python -m unittest discover -s tests/inference -p "test_*.py"

pyproject.toml

@@ -120,7 +120,8 @@ inference-tests = [
   "litellm==v1.52.9",
   "tenacity",
   "diskcache",
-  "numpy==1.26.4"
+  "numpy==1.26.4",
+  "ollama"
 ]
 assistant = [
     "streamlit",

src/unitxt/inference.py

@@ -1156,15 +1156,14 @@ def _infer(
         import ollama
 
         args = self.to_dict([StandardAPIParamsMixin])
-
         results = []
-
+        model = args.pop("model")
         for instance in dataset:
             messages = self.to_messages(instance)
             response = ollama.chat(
-                model=self.model,
                 messages=messages,
-                **args,
+                model=model,
+                options=args,
             )
             results.append(response)
 

tests/inference/test_inference_engine.py

@@ -9,6 +9,7 @@
     HFOptionSelectingInferenceEngine,
     HFPipelineBasedInferenceEngine,
     LiteLLMInferenceEngine,
+    OllamaInferenceEngine,
     OptionSelectingByLogProbsInferenceEngine,
     RITSInferenceEngine,
     TextGenerationInferenceOutput,
@@ -402,3 +403,13 @@ def test_hugginface_pipeline_inference_engine_chat_api(self):
 
         self.assertEqual(predictions[0], "Hello")
         self.assertEqual(predictions[1], "As")
+
+    def test_ollama_inference_engine(self):
+        dataset = [
+            {"source": "Answer in one word only. What is the capital of Canada"},
+        ]
+
+        engine = OllamaInferenceEngine(model="llama3.2:1b", temperature=0.0)
+        predictions = engine.infer(dataset)
+
+        self.assertTrue("Ottawa" in predictions[0], predictions[0])

utils/.secrets.baseline

@@ -133,14 +133,14 @@
         "filename": "src/unitxt/inference.py",
         "hashed_secret": "aa6cd2a77de22303be80e1f632195d62d211a729",
         "is_verified": false,
-        "line_number": 1294
+        "line_number": 1293
       },
       {
         "type": "Secret Keyword",
         "filename": "src/unitxt/inference.py",
         "hashed_secret": "c8f16a194efc59559549c7bd69f7bea038742e79",
         "is_verified": false,
-        "line_number": 1779
+        "line_number": 1778
       }
     ],
     "src/unitxt/loaders.py": [
@@ -178,5 +178,5 @@
       }
     ]
   },
-  "generated_at": "2025-02-17T14:14:28Z"
+  "generated_at": "2025-02-17T19:31:06Z"
 }