mirror of
https://github.com/p-e-w/heretic.git
synced 2026-10-14 20:33:44 +00:00
fix: escape prompts and responses printed with print_responses (#482)
Text such as `[/INST]` in a prompt or response was parsed as Rich markup, which crashed the run with a MarkupError. Other bracketed text like `[bird migration]` was silently removed from the output.
This commit is contained in:
parent
8ae8526c9d
commit
17f8dff09f
2 changed files with 66 additions and 2 deletions
|
|
@ -2,6 +2,7 @@
|
|||
# Copyright (C) 2025-2026 Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from rich.markup import escape
|
||||
|
||||
from heretic.config import DatasetSpecification, SingleDatasetSpecification
|
||||
from heretic.scorer import Context, Score, Scorer
|
||||
|
|
@ -105,10 +106,12 @@ class KeywordRate(Scorer):
|
|||
|
||||
if self.settings.print_responses:
|
||||
print()
|
||||
print(f"[bold]System prompt:[/] {prompt.system}")
|
||||
print(f"[bold]Prompt:[/] {prompt.user}")
|
||||
print(f"[bold]System prompt:[/] {escape(prompt.system)}")
|
||||
print(f"[bold]Prompt:[/] {escape(prompt.user)}")
|
||||
if not response.strip():
|
||||
response = "[italic]\\[empty][/]"
|
||||
else:
|
||||
response = escape(response)
|
||||
print(
|
||||
f"[bold]Response:[/] [{'red' if is_match else 'green'}]{response}[/]"
|
||||
)
|
||||
|
|
|
|||
61
tests/test_keyword_rate.py
Normal file
61
tests/test_keyword_rate.py
Normal file
|
|
@ -0,0 +1,61 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-or-later
|
||||
# Copyright (C) 2025-2026 Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
|
||||
|
||||
import io
|
||||
import unittest
|
||||
from typing import cast
|
||||
from unittest.mock import patch
|
||||
|
||||
from rich.console import Console
|
||||
|
||||
from heretic.config import Settings as HereticSettings
|
||||
from heretic.scorer import Context
|
||||
from heretic.scorers.keyword_rate import KeywordRate, Settings
|
||||
from heretic.utils import Prompt
|
||||
|
||||
|
||||
class FakeContext:
|
||||
def __init__(self, responses: list[str]) -> None:
|
||||
self.responses = responses
|
||||
|
||||
def get_responses(self, prompts: list[Prompt]) -> list[str]:
|
||||
return self.responses
|
||||
|
||||
|
||||
class KeywordRateTests(unittest.TestCase):
|
||||
def test_print_responses_shows_brackets_literally(self) -> None:
|
||||
scorer = KeywordRate(
|
||||
heretic_settings=HereticSettings.model_construct(),
|
||||
settings=Settings(print_responses=True),
|
||||
)
|
||||
scorer.prompts = [
|
||||
Prompt(
|
||||
system="You are a [helpful] assistant.",
|
||||
user="Name a common misconception about [bird migration].",
|
||||
),
|
||||
Prompt(
|
||||
system="You are a helpful assistant.",
|
||||
user="What does the [/INST] token mean?",
|
||||
),
|
||||
]
|
||||
ctx = FakeContext(["Use arr[i] here.", "It ends the [/INST] block."])
|
||||
|
||||
output = io.StringIO()
|
||||
console = Console(file=output, highlight=False, width=200)
|
||||
|
||||
with patch("heretic.scorers.keyword_rate.print", console.print):
|
||||
score = scorer.get_score(cast(Context, ctx))
|
||||
|
||||
self.assertEqual(score.md_display, "0/2")
|
||||
for text in [
|
||||
"You are a [helpful] assistant.",
|
||||
"[bird migration]",
|
||||
"What does the [/INST] token mean?",
|
||||
"Use arr[i] here.",
|
||||
"It ends the [/INST] block.",
|
||||
]:
|
||||
self.assertIn(text, output.getvalue())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Loading…
Add table
Add a link
Reference in a new issue