hallucination with log probs (#281)

* first init * fix * fix test * new implemenetation * fix bug * fix bug * fix bug * address issue * address issues * address comments * fix test * fix * move constatns * remove consts
2026-04-28 18:36:34 +02:00 · 2024-11-27 15:17:02 -08:00 · 2024-11-27 15:17:02 -08:00 · cadd3cdaf9
commit cadd3cdaf9
parent f5cdafb7c8
5 changed files with 1269 additions and 1 deletions
--- a/model_server/app/tests/test_hallucination.py
+++ b/model_server/app/tests/test_hallucination.py
@ -0,0 +1,148 @@
+import json
+from app.function_calling.hallucination_handler import HallucinationStateHandler
+import pytest
+import os
+
+# Get the directory of the current file
+current_dir = os.path.dirname(__file__)
+
+# Construct the full path to the JSON file
+json_file_path = os.path.join(current_dir, "test_cases.json")
+
+with open(json_file_path) as f:
+    test_cases = json.load(f)
+
+get_weather_api = {
+    "type": "function",
+    "function": {
+        "name": "get_current_weather",
+        "description": "Get current weather at a location.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "location": {
+                    "type": "str",
+                    "description": "The location to get the weather for",
+                    "format": "City, State",
+                },
+                "unit": {
+                    "type": "str",
+                    "description": "The unit to return the weather in.",
+                    "enum": ["celsius", "fahrenheit"],
+                    "default": "celsius",
+                },
+                "days": {
+                    "type": "str",
+                    "description": "the number of days for the request.",
+                },
+            },
+            "required": ["location", "days"],
+        },
+    },
+}
+function_description = get_weather_api["function"]
+if type(function_description) != list:
+    function_description = [get_weather_api["function"]]
+
+
+@pytest.mark.parametrize("case", test_cases)
+def test_hallucination(case):
+    state = HallucinationStateHandler(
+        response_iterator=None, function=function_description
+    )
+    for token, logprob in zip(case["tokens"], case["logprobs"]):
+        if token != "</tool_call>":
+            state.append_and_check_token_hallucination(token, logprob)
+            if state.hallucination:
+                break
+    assert state.hallucination == case["expect"]
+
+
+@pytest.mark.parametrize("is_hallucinate_sample", [True, False])
+def test_hallucination_prompt(is_hallucinate_sample):
+    TASK_PROMPT = """
+    You are a helpful assistant.
+    """.strip()
+
+    TOOL_PROMPT = """
+    # Tools
+
+    You may call one or more functions to assist with the user query.
+
+    You are provided with function signatures within <tools></tools> XML tags:
+    <tools>
+    {tool_text}
+    </tools>
+    """.strip()
+
+    FORMAT_PROMPT = """
+    For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
+    <tool_call>
+    {"name": <function-name>, "arguments": <args-json-object>}
+    </tool_call>
+    """.strip()
+
+    def convert_tools(tools):
+        return "\n".join([json.dumps(tool) for tool in tools])
+
+    def format_prompt(tools):
+        tool_text = convert_tools(tools)
+
+        return (
+            TASK_PROMPT
+            + "\n\n"
+            + TOOL_PROMPT.format(tool_text=tool_text)
+            + "\n\n"
+            + FORMAT_PROMPT
+            + "\n"
+        )
+
+    openai_format_tools = [get_weather_api]
+
+    system_prompt = format_prompt(openai_format_tools)
+
+    from openai import OpenAI
+
+    client = OpenAI(base_url="https://api.fc.archgw.com/v1", api_key="EMPTY")
+
+    # List models API
+    model = client.models.list().data[0].id
+    assert model == "Arch-Function"
+    if not is_hallucinate_sample:
+        messages = [
+            {"role": "system", "content": system_prompt},
+            # {"role": "user", "content": "can you help me check weather?"},
+            {"role": "user", "content": "How is the weather in Seattle in 7 days?"},
+            # {"role": "assistant", "content": "Of course!"},
+            # {"role": "user", "content": "Seattle please"}
+        ]
+    else:
+        messages = [
+            {"role": "system", "content": system_prompt},
+            # {"role": "user", "content": "can you help me check weather?"},
+            {"role": "user", "content": "How is the weather in Seattle in days?"},
+            # {"role": "assistant", "content": "Of course!"},
+            # {"role": "user", "content": "Seattle please"}
+        ]
+
+    extra_body = {
+        "temperature": 0.6,
+        "top_p": 1.0,
+        "top_k": 50,
+        # "continue_final_message": True,
+        # "add_generation_prompt": False,
+        "logprobs": True,
+        "top_logprobs": 10,
+    }
+
+    resp = client.chat.completions.create(
+        model="Arch-Function", messages=messages, extra_body=extra_body, stream=True
+    )
+
+    hallu = HallucinationStateHandler(
+        response_iterator=resp, function=function_description
+    )
+
+    for token in hallu:
+        assert len(hallu.tokens) >= 0
+    assert hallu.hallucination == is_hallucinate_sample