src/axolotl/prompt_strategies/jinja_template_analyzer.py (2)

32-34: Consume the tag token explicitly for clarity

parser.stream.skip(1) silently advances the stream but obscures intent. Using next(parser.stream) (as in Jinja’s own examples) makes it obvious that the tag token is consumed and conveys the current lineno for potential error reporting.

-    def parse(self, parser):
-        parser.stream.skip(1)
-        return nodes.Const("")
+    def parse(self, parser):
+        # Discard the tag token and produce an empty node
+        next(parser.stream)          # consume 'generation'
+        return nodes.Const("")

---

66-66: Offer an extension-injection hook

Hard-coding extensions=[GenerationTagIgnore] makes the analyser less composable for downstream users who might already be passing custom extensions (e.g., i18n or sandbox). Accepting an optional extra_extensions parameter and merging it would avoid forcing callers to fork this class.

-    def __init__(self, template: str):
-        self.env: Environment = Environment(autoescape=True, extensions=[GenerationTagIgnore])
+    def __init__(self, template: str, *, extra_extensions: Optional[list[type[Extension]]] = None):
+        exts = [GenerationTagIgnore]
+        if extra_extensions:
+            exts.extend(extra_extensions)
+        self.env: Environment = Environment(autoescape=True, extensions=exts)

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

src/axolotl/utils/chat_templates.py (1)

35-35: phi_4 template – edge-cases & maintainability review

The template works but a couple of minor issues surfaced:

1. New-line consistency – the <|im_start|>assistant<|im_sep|> token is concatenated directly with the following {% generation %} block. Lack of a newline or space may yield <|im_sep|><think> if the assistant’s content starts with <think>. Consider adding an explicit '\n' after the separator to ensure clean token boundaries.

2. Unhandled roles – the loop omits handling for a second system or a future tool role. That might be intentional, but it differs from other built-in templates. If unexpected roles can appear, add an else branch that raises to aid debugging.

3. Readability – the gigantic default system prompt makes the template hard to scan. Extracting it into a dedicated variable keeps the Jinja logic short:

-{% else %}You are Phi, a language model ...
+{% else %}
+{% set default_phi_prompt %}
+You are Phi, a language model ...
+(unchanged text)
+{% endset %}
+{{ default_phi_prompt }}
{% endif %}

(Not blocking, but will help future maintenance.)

tests/prompt_strategies/test_chat_templates_advanced.py (3)

36-44: Parameter list keeps ballooning – extract to data fixture to avoid duplication

We now hard-code 9-tuple entries; each new tokenizer / template combination repeats 4 strings.
Maintaining this inline list is error-prone (ordering, duplicates, forgotten cleanup).
Consider moving the matrix to a small JSON / YAML or a pytest fixture that yields the tuples, e.g.

# conftest.py
CHAT_TEMPLATE_PARAMS = [
    ("llama3_tokenizer", "llama3", None, None),
    ...
]

@pytest.fixture(params=CHAT_TEMPLATE_PARAMS, ids=lambda p: p[0])
def chat_template_case(request):
    return request.param

and update the test with @pytest.mark.parametrize(*chat_template_case.__annotations__.keys()).

This keeps the test readable and makes additions like gemma2_tokenizer / phi4_tokenizer one-liners instead of 4-liners scattered in the middle of a giant list.

---

95-98: _should_skip_turn hard-codes tokenizer name – make it declarative

The skip condition is now limited to "mistral" tokenizers, but the string test is embedded in the method:

and ("mistral" in tokenizer.name_or_path.lower())

Encapsulate the knowledge of “which tokenizers elide the first system/context turn” in a small set for maintainability:

+_SYSTEM_SKIP_TOKENIZERS = {"mistral"}

...
-            and ("mistral" in tokenizer.name_or_path.lower())
+            and any(t in tokenizer.name_or_path.lower() for t in _SYSTEM_SKIP_TOKENIZERS)

Adding future models becomes a one-line update to the set instead of a code change.

---

963-968: Duplicate elif branches – collapse into a single membership check

These two branches perform the identical assertion; Ruff rightly flags SIM114.
You can combine them for brevity and to avoid future copy-paste errors:

-        elif chat_template == "phi_35":
-            assert variables == {"role", "content"}, (
-                f"Expected variables: {'role', 'content'} from {tokenizer}/{chat_template}\n"
-                f"Got: {variables}\n"
-                f"Chat template: {actual_jinja_template}"
-            )
-        elif chat_template == "phi_4":
+        elif chat_template in {"phi_35", "phi_4"}:
             assert variables == {"role", "content"}, (
                 f"Expected variables: {'role', 'content'} from {tokenizer}/{chat_template}\n"
                 f"Got: {variables}\n"
                 f"Chat template: {actual_jinja_template}"
             )

This removes redundant code and satisfies the static-analysis suggestion.

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro