test_ggml.py

import unittest
from unittest.mock import MagicMock
from lm_eval.models.ggml import GGMLLM

class GGMLLMTest(unittest.TestCase):
    def test_loglikelihood(self):
        base_url = "https://matthoffner-ggml-llm-api.hf.space"
        lm = GGMLLM(base_url)

        # Create a MagicMock object to mock ggml_completion
        ggml_completion_mock = MagicMock()

        # Set the return value for the mocked function
        ggml_completion_mock.return_value = {
            "logprob": -1.2345,
            "is_greedy": True
        }

        # Patch the ggml_completion function with the mocked function
        lm.ggml_completion = ggml_completion_mock

        # Test loglikelihood
        requests = [("context1", "continuation1"), ("context2", "continuation2")]
        res = lm.loglikelihood(requests)

        # Assert the loglikelihood response is correct
        expected_res = [(-1.2345, True), (-1.2345, True)]
        self.assertEqual(res, expected_res)

    def test_greedy_until(self):
        base_url = "https://matthoffner-ggml-llm-api.hf.space"
        lm = GGMLLM(base_url)

        # Define the ggml_completion method with the desired behavior
        def ggml_completion_mock(url, context, stop=None):
            if stop is not None:
                return {"text": f"generated_text{stop[-1]}"}
            return {"text": "generated_text"}

        # Set the ggml_completion method to the defined mock
        lm.ggml_completion = ggml_completion_mock

        # Test greedy_until
        requests = [("input1", {"until": "stop1"}), ("input2", {"until": "stop2"})]
        res = lm.greedy_until(requests)

        # Assert the greedy_until response is correct
        expected_res = ["generated_text1", "generated_text2"]
        self.assertEqual(res, expected_res)

if __name__ == "__main__":
    unittest.main()