Repository navigation
Feature/chatml schemas - #2
Conversation
📝 WalkthroughWalkthroughThe PR rewrites the chat completions API to use Hugging Face chat templates, sampling controls, and stopping criteria; expands request/response schemas with strict validation; refactors application settings to typed lowercase fields; adds a legacy-to-ChatML data migration script; and restructures training into a class-based QLoRA pipeline. ChangesChatML API, schema, config, migration, and trainer refactor
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant ChatEndpoint
participant Tokenizer
participant Model
Client->>ChatEndpoint: POST /completions (GenerationRequest)
ChatEndpoint->>Tokenizer: apply_chat_template(messages)
Tokenizer-->>ChatEndpoint: prompt_text, input_ids
ChatEndpoint->>Model: generate(input_ids, stopping_criteria, sampling params)
Model-->>ChatEndpoint: generated token ids
ChatEndpoint->>Tokenizer: decode new tokens
ChatEndpoint-->>Client: GenerationResponse (text, tokens, finish_reason)
sequenceDiagram
participant Train
participant LLMForgeTrainer
participant Tokenizer
participant Model
participant SFTTrainer
Train->>LLMForgeTrainer: train()
LLMForgeTrainer->>LLMForgeTrainer: load_dataset()
LLMForgeTrainer->>Tokenizer: configure ChatML template
LLMForgeTrainer->>Model: load quantized base model
LLMForgeTrainer->>SFTTrainer: construct(tokenizer, args, peft_config)
SFTTrainer-->>LLMForgeTrainer: trained model
LLMForgeTrainer-->>Train: save model/tokenizer to adapter_path
Related issues: None specified. Related PRs: None specified. Suggested labels: enhancement, api, training Suggested reviewers: Chanakya-22 🐰 A chat that once rambled now speaks with care, 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 12
🧹 Nitpick comments (1)
src/train.py (1)
129-153: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winConsider enabling gradient checkpointing for QLoRA.
Given the explicit
torch.OutOfMemoryErrorhandling and constrained batch size (per_device_train_batch_size=1,gradient_accumulation_steps=4), enablinggradient_checkpointing=TrueinSFTConfig(paired withmodel.config.use_cache = False, already set) is a standard QLoRA memory-saving practice worth adding.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/train.py` around lines 129 - 153, The training config in _build_training_config should enable gradient checkpointing to reduce memory usage for QLoRA. Update the SFTConfig returned by _build_training_config to set gradient_checkpointing=True, and make sure the model setup keeps use_cache disabled as already handled so checkpointing works correctly with the existing low batch size and OutOfMemoryError handling.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/app/api/v1/chat.py`:
- Around line 135-136: The finish_reason logic in the chat generation path is
using only tokens_generated versus request.max_new_tokens, which mislabels a
normal EOS/stop completion as "length" when it happens exactly at the limit.
Update the finish-reason calculation in the chat endpoint logic around
tokens_generated to inspect the final generated token/sequence for EOS or stop
conditions first, and only use "length" when generation truly ended due to the
token cap.
- Around line 34-58: The current stop logic in the stop criteria class matches
any single token ID, which can terminate generation too early for multi-token
stop sequences. Update the stop handling in the relevant stop criteria used by
chat generation so it checks the full trailing token subsequence against each
configured stop sequence instead of only comparing the latest token; use the
existing stop criteria class and its __call__ method as the location to adjust,
and ensure the same fix is applied to the other referenced stop-criteria block
as well.
- Around line 153-160: The generic exception handler in chat.py is exposing raw
exception text via JSONResponse, which should not be sent to clients. Update the
broad except Exception as exc block to return only a sanitized error message in
the response, and log the full exception details server-side instead. Use the
existing JSONResponse handling in the chat endpoint to keep the client-facing
payload generic while preserving debuggability through server logs.
In `@src/app/schemas/protocol.py`:
- Around line 76-80: The unbounded user-controlled arrays in the protocol schema
can drive excessive templating and tokenization work in the chat flow. Add
explicit maximum-size validation for messages and stop_sequences in the Protocol
schema, and make sure the chat handling in chat.py respects those limits before
calling apply_chat_template, tokenizer, or per-stop-sequence encoding. Use the
existing messages and stop_sequences fields as the primary points to enforce the
guardrails.
- Around line 52-53: The string-type check in the protocol schema currently
raises TypeError, which escapes Pydantic validation instead of becoming a
ValidationError. Update the validation logic in the relevant validator that
checks message content to raise ValueError or PydanticCustomError instead of
TypeError, so invalid non-string input is reported through Pydantic’s normal
validation flow.
- Around line 87-98: The sampling fields in Protocol schema currently allow zero
values, which can break `model.generate(...)` when `do_sample` is enabled.
Update the validation around the `temperature` and `top_p` settings in
`src/app/schemas/protocol.py` so that zero is rejected only for sampling
requests, while greedy mode can still ignore these parameters. Use the existing
`do_sample` handling in `Protocol`/request validation to enforce strictly
positive values for `temperature` and `top_p` when sampling is on.
In `@src/core/config.py`:
- Around line 44-47: The optional default on adapter_path in the settings model
conflicts with LLMForgeTrainer.__init__ in train.py, which unconditionally
passes settings.adapter_path into Path() when building self.output_dir. Make
adapter_path required for training or introduce a separate required output_dir
setting, and add validation in the settings/trainer initialization so the
trainer fails fast with a clear error before Path() is called on a missing
value.
- Around line 40-43: `load_model()` in `src/engine/generator.py` is still using
removed uppercase settings names, so it will fail with `AttributeError` because
`Settings` now exposes `base_model_name` and `adapter_path`. Update the
`load_model` call site to read the lowercase `settings.base_model_name` and
`settings.adapter_path`, or add explicit backward-compatible aliases on
`Settings` if those uppercase names must remain supported.
In `@src/core/migrate_data.py`:
- Around line 27-33: The legacy record parsing in migrate_data.py is silently
truncating multi-turn conversations by splitting on "<bot>:" and using only the
first split result. Update the parsing logic in the migration loop to extract
only the last assistant turn from raw_text (for example, via rsplit on
"<bot>:"), and add validation/error handling so records with extra turns are not
merged or dropped without notice. Ensure the transformation around user_part and
assistant_part preserves the intended final turn only.
In `@src/train.py`:
- Around line 71-74: The validation split produced in _load_dataset is currently
only logged and then discarded, while _build_trainer still passes
eval_dataset=None even though _build_training_config sets eval_strategy to
steps. Update _load_dataset to retain or return both the train and test splits,
then wire the saved validation split into _build_trainer as the eval_dataset so
Trainer/SFTTrainer can evaluate when eval_steps is reached.
- Around line 43-44: The trainer initialization in LLMForgeTrainer is crashing
because output_dir is built from settings.adapter_path even when that value is
None. Update the constructor logic around self.output_dir so it handles an unset
adapter_path safely, either by deferring Path conversion until a real path is
provided or by falling back to a sensible default before instantiating Path.
Keep the fix localized to the LLMForgeTrainer constructor and any related config
access in Settings.adapter_path.
- Around line 186-193: The SFTTrainer construction in train.py is still passing
tokenizer=, but newer unpinned trl releases have removed that argument. Update
the SFTTrainer call in the trainer-building method to use processing_class= with
the same tokenizer/processor object instead, and keep the rest of the arguments
unchanged so it works across current releases.
---
Nitpick comments:
In `@src/train.py`:
- Around line 129-153: The training config in _build_training_config should
enable gradient checkpointing to reduce memory usage for QLoRA. Update the
SFTConfig returned by _build_training_config to set gradient_checkpointing=True,
and make sure the model setup keeps use_cache disabled as already handled so
checkpointing works correctly with the existing low batch size and
OutOfMemoryError handling.
🪄 Autofix (Beta)
✅ Autofix completed
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 408bd6c3-2747-4922-9d6a-2d72d62980ad
📒 Files selected for processing (5)
src/app/api/v1/chat.pysrc/app/schemas/protocol.pysrc/core/config.pysrc/core/migrate_data.pysrc/train.py
| def __init__(self, stop_token_ids: Optional[List[int]] = None) -> None: | ||
| """Initialize the criteria with explicit token IDs to stop on. | ||
|
|
||
| Args: | ||
| stop_token_ids: A list of token IDs that should terminate generation. | ||
| """ | ||
| super().__init__() | ||
| self.stop_token_ids = set(stop_token_ids or []) | ||
|
|
||
| def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **_: object) -> bool: | ||
| """Check whether generation should stop at the current token step. | ||
|
|
||
| Args: | ||
| input_ids: The full token sequence produced so far. | ||
| scores: The logits for the current token step. | ||
|
|
||
| Returns: | ||
| ``True`` if the latest generated token is a stop token, otherwise | ||
| ``False``. | ||
| """ | ||
| if not self.stop_token_ids: | ||
| return False | ||
|
|
||
| last_token_id = int(input_ids[0, -1].item()) | ||
| return last_token_id in self.stop_token_ids |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Match complete stop sequences, not individual token IDs.
The current union of token IDs stops generation when any token from a stop sequence appears. Multi-token stop sequences can therefore terminate on a partial/common token before the full stop marker is produced.
Proposed fix
- def __init__(self, stop_token_ids: Optional[List[int]] = None) -> None:
+ def __init__(self, stop_sequences: Optional[List[List[int]]] = None) -> None:
@@
- self.stop_token_ids = set(stop_token_ids or [])
+ self.stop_sequences = [sequence for sequence in (stop_sequences or []) if sequence]
@@
- if not self.stop_token_ids:
+ if not self.stop_sequences:
return False
- last_token_id = int(input_ids[0, -1].item())
- return last_token_id in self.stop_token_ids
+ generated_ids = input_ids[0].tolist()
+ return any(
+ len(generated_ids) >= len(sequence)
+ and generated_ids[-len(sequence):] == sequence
+ for sequence in self.stop_sequences
+ )
@@
- stop_token_ids = set()
+ stop_token_sequences = []
for stop_sequence in request.stop_sequences:
if not stop_sequence:
continue
token_ids = tokenizer.encode(stop_sequence, add_special_tokens=False)
if token_ids:
- stop_token_ids.update(token_ids)
+ stop_token_sequences.append(token_ids)
if getattr(tokenizer, "eos_token_id", None) is not None:
- stop_token_ids.add(int(tokenizer.eos_token_id))
+ stop_token_sequences.append([int(tokenizer.eos_token_id)])
stopping_criteria = StoppingCriteriaList(
- [TokenStopCriteria(list(stop_token_ids))]
+ [TokenStopCriteria(stop_token_sequences)]
)Also applies to: 93-106
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/app/api/v1/chat.py` around lines 34 - 58, The current stop logic in the
stop criteria class matches any single token ID, which can terminate generation
too early for multi-token stop sequences. Update the stop handling in the
relevant stop criteria used by chat generation so it checks the full trailing
token subsequence against each configured stop sequence instead of only
comparing the latest token; use the existing stop criteria class and its
__call__ method as the location to adjust, and ensure the same fix is applied to
the other referenced stop-criteria block as well.
| tokens_generated = int(generated_token_ids.shape[0]) | ||
| finish_reason = "stop" if tokens_generated < request.max_new_tokens else "length" |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Base finish_reason on the actual stop condition.
If EOS/stop is emitted exactly at max_new_tokens, this reports "length" even though generation stopped normally. Check the final generated token/sequence before falling back to "length".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/app/api/v1/chat.py` around lines 135 - 136, The finish_reason logic in
the chat generation path is using only tokens_generated versus
request.max_new_tokens, which mislabels a normal EOS/stop completion as "length"
when it happens exactly at the limit. Update the finish-reason calculation in
the chat endpoint logic around tokens_generated to inspect the final generated
token/sequence for EOS or stop conditions first, and only use "length" when
generation truly ended due to the token cap.
| except Exception as exc: | ||
| return JSONResponse( | ||
| status_code=500, | ||
| content={ | ||
| "detail": "Generation failed due to an internal server error.", | ||
| "error": str(exc), | ||
| }, | ||
| ) No newline at end of file |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
Don’t return raw exception strings to clients.
The generic handler exposes str(exc) and catches all exceptions. Return a sanitized message and log details server-side instead.
Minimal response hardening
except Exception as exc:
+ # TODO: log exception details server-side.
return JSONResponse(
status_code=500,
content={
"detail": "Generation failed due to an internal server error.",
- "error": str(exc),
+ "error": "internal_error",
},
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| except Exception as exc: | |
| return JSONResponse( | |
| status_code=500, | |
| content={ | |
| "detail": "Generation failed due to an internal server error.", | |
| "error": str(exc), | |
| }, | |
| ) | |
| except Exception as exc: | |
| # TODO: log exception details server-side. | |
| return JSONResponse( | |
| status_code=500, | |
| content={ | |
| "detail": "Generation failed due to an internal server error.", | |
| "error": "internal_error", | |
| }, | |
| ) |
🧰 Tools
🪛 Ruff (0.15.20)
[warning] 153-153: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/app/api/v1/chat.py` around lines 153 - 160, The generic exception handler
in chat.py is exposing raw exception text via JSONResponse, which should not be
sent to clients. Update the broad except Exception as exc block to return only a
sanitized error message in the response, and log the full exception details
server-side instead. Use the existing JSONResponse handling in the chat endpoint
to keep the client-facing payload generic while preserving debuggability through
server logs.
Source: Linters/SAST tools
| if not isinstance(value, str): | ||
| raise TypeError("Message content must be provided as a string.") |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python - <<'PY'
from pydantic import BaseModel, field_validator
class M(BaseModel):
content: str
`@field_validator`("content", mode="before")
`@classmethod`
def validate_content(cls, value):
if not isinstance(value, str):
raise TypeError("Message content must be provided as a string.")
return value
try:
M.model_validate({"content": 123})
except Exception as exc:
print(type(exc).__name__)
print(exc)
PYRepository: Chanakya-22/mini_llm_forge
Length of output: 220
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== file outline ==\n'
ast-grep outline src/app/schemas/protocol.py --view expanded || true
printf '\n== relevant lines ==\n'
sed -n '1,140p' src/app/schemas/protocol.py | cat -n
printf '\n== pydantic version ==\n'
python3 - <<'PY'
import importlib.metadata as md
try:
print(md.version("pydantic"))
except Exception as e:
print(type(e).__name__, e)
PYRepository: Chanakya-22/mini_llm_forge
Length of output: 6921
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pydantic import BaseModel, field_validator
class MType(BaseModel):
content: str
`@field_validator`("content", mode="before")
`@classmethod`
def validate_content(cls, value):
if not isinstance(value, str):
raise TypeError("Message content must be provided as a string.")
return value
class MValue(BaseModel):
content: str
`@field_validator`("content", mode="before")
`@classmethod`
def validate_content(cls, value):
if not isinstance(value, str):
raise ValueError("Message content must be provided as a string.")
return value
for cls in (MType, MValue):
try:
cls.model_validate({"content": 123})
except Exception as exc:
print(cls.__name__, type(exc).__name__, str(exc).splitlines()[0])
PYRepository: Chanakya-22/mini_llm_forge
Length of output: 279
Return a validation error here instead of TypeError. Raise ValueError (or PydanticCustomError) for invalid input so Pydantic turns this into a ValidationError instead of letting the exception escape.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/app/schemas/protocol.py` around lines 52 - 53, The string-type check in
the protocol schema currently raises TypeError, which escapes Pydantic
validation instead of becoming a ValidationError. Update the validation logic in
the relevant validator that checks message content to raise ValueError or
PydanticCustomError instead of TypeError, so invalid non-string input is
reported through Pydantic’s normal validation flow.
| messages: List[ChatMessage] = Field( | ||
| ..., | ||
| min_length=1, | ||
| description="An ordered list of ChatML messages representing the conversation history.", | ||
| ) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Cap array sizes before they reach templating and tokenization.
messages and stop_sequences are user-controlled and unbounded. A large list can force excessive apply_chat_template, tokenizer, and per-stop-sequence encoding work in src/app/api/v1/chat.py.
Proposed guardrails
messages: List[ChatMessage] = Field(
...,
min_length=1,
+ max_length=128,
description="An ordered list of ChatML messages representing the conversation history.",
)
@@
stop_sequences: List[str] = Field(
default_factory=lambda: ["<|im_end|>"],
+ max_length=16,
description="A list of stop sequences that terminate generation when encountered.",
)Also applies to: 115-118
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/app/schemas/protocol.py` around lines 76 - 80, The unbounded
user-controlled arrays in the protocol schema can drive excessive templating and
tokenization work in the chat flow. Add explicit maximum-size validation for
messages and stop_sequences in the Protocol schema, and make sure the chat
handling in chat.py respects those limits before calling apply_chat_template,
tokenizer, or per-stop-sequence encoding. Use the existing messages and
stop_sequences fields as the primary points to enforce the guardrails.
| adapter_path: Optional[str] = Field( | ||
| default=None, | ||
| description="Optional filesystem path to a LoRA adapter directory.", | ||
| ) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
adapter_path defaulting to None is incompatible with the trainer's required usage.
src/train.py's LLMForgeTrainer.__init__ does Path(settings.adapter_path) unconditionally to build self.output_dir. With this field optional and defaulting to None, instantiating the trainer without setting ADAPTER_PATH/adapter_path in the environment will raise TypeError: argument should be a str or PathLike, not NoneType immediately.
Consider making adapter_path required for training (or add a distinct output_dir setting) and validate it before use.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/core/config.py` around lines 44 - 47, The optional default on
adapter_path in the settings model conflicts with LLMForgeTrainer.__init__ in
train.py, which unconditionally passes settings.adapter_path into Path() when
building self.output_dir. Make adapter_path required for training or introduce a
separate required output_dir setting, and add validation in the settings/trainer
initialization so the trainer fails fast with a clear error before Path() is
called on a missing value.
| if "<human>:" not in raw_text or "<bot>:" not in raw_text: | ||
| logging.warning(f"Skipping line {line_num}: Invalid legacy format.") | ||
| continue | ||
|
|
||
| parts = raw_text.split("<bot>:") | ||
| user_part = parts[0].replace("<human>:", "").strip() | ||
| assistant_part = parts[1].strip() |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Multi-turn legacy records are silently corrupted.
Splitting on "<bot>:" and taking only parts[1] assumes exactly one <bot>: marker. If a legacy record has multiple turns (e.g. <human>: Q1 <bot>: A1 <human>: Q2 <bot>: A2), this will merge the next human turn into the assistant response and drop the final turn entirely, silently producing corrupted training data instead of erroring out.
🛠️ Proposed fix using rsplit for last-turn extraction
- parts = raw_text.split("<bot>:")
- user_part = parts[0].replace("<human>:", "").strip()
- assistant_part = parts[1].strip()
+ user_part, _, assistant_part = raw_text.partition("<bot>:")
+ user_part = user_part.replace("<human>:", "").strip()
+ assistant_part = assistant_part.strip()
+ if "<bot>:" in assistant_part or "<human>:" in assistant_part:
+ logging.warning(f"Skipping line {line_num}: Multi-turn record not supported.")
+ continue📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if "<human>:" not in raw_text or "<bot>:" not in raw_text: | |
| logging.warning(f"Skipping line {line_num}: Invalid legacy format.") | |
| continue | |
| parts = raw_text.split("<bot>:") | |
| user_part = parts[0].replace("<human>:", "").strip() | |
| assistant_part = parts[1].strip() | |
| if "<human>:" not in raw_text or "<bot>:" not in raw_text: | |
| logging.warning(f"Skipping line {line_num}: Invalid legacy format.") | |
| continue | |
| user_part, _, assistant_part = raw_text.partition("<bot>:") | |
| user_part = user_part.replace("<human>:", "").strip() | |
| assistant_part = assistant_part.strip() | |
| if "<bot>:" in assistant_part or "<human>:" in assistant_part: | |
| logging.warning(f"Skipping line {line_num}: Multi-turn record not supported.") | |
| continue |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/core/migrate_data.py` around lines 27 - 33, The legacy record parsing in
migrate_data.py is silently truncating multi-turn conversations by splitting on
"<bot>:" and using only the first split result. Update the parsing logic in the
migration loop to extract only the last assistant turn from raw_text (for
example, via rsplit on "<bot>:"), and add validation/error handling so records
with extra turns are not merged or dropped without notice. Ensure the
transformation around user_part and assistant_part preserves the intended final
turn only.
| self.dataset_path = Path(dataset_path) | ||
| self.output_dir = Path(settings.adapter_path) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
Trainer construction crashes when adapter_path is unset.
Path(settings.adapter_path) is called unconditionally, but Settings.adapter_path defaults to None (see src/core/config.py). Simply instantiating LLMForgeTrainer() without configuring adapter_path will raise a TypeError before any training logic runs.
🛠️ Proposed fix
def __init__(self, dataset_path: str = "data/chatml_data.jsonl") -> None:
self.dataset_path = Path(dataset_path)
- self.output_dir = Path(settings.adapter_path)
+ if not settings.adapter_path:
+ raise ValueError("settings.adapter_path must be configured for training.")
+ self.output_dir = Path(settings.adapter_path)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| self.dataset_path = Path(dataset_path) | |
| self.output_dir = Path(settings.adapter_path) | |
| self.dataset_path = Path(dataset_path) | |
| if not settings.adapter_path: | |
| raise ValueError("settings.adapter_path must be configured for training.") | |
| self.output_dir = Path(settings.adapter_path) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/train.py` around lines 43 - 44, The trainer initialization in
LLMForgeTrainer is crashing because output_dir is built from
settings.adapter_path even when that value is None. Update the constructor logic
around self.output_dir so it handles an unset adapter_path safely, either by
deferring Path conversion until a real path is provided or by falling back to a
sensible default before instantiating Path. Keep the fix localized to the
LLMForgeTrainer constructor and any related config access in
Settings.adapter_path.
| split_dataset = dataset.train_test_split(test_size=0.1, seed=42) | ||
| self.dataset = split_dataset["train"] | ||
| logger.info("Loaded %s training examples and %s validation examples.", len(self.dataset), len(split_dataset["test"])) | ||
| return self.dataset |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
Validation split is discarded; eval_dataset=None conflicts with eval_strategy="steps".
_load_dataset computes split_dataset["test"] split_dataset = dataset.train_test_split(test_size=0.1, seed=42) but only logs its length and never stores or returns it — self.dataset is set solely to the train split. _build_trainer then hardcodes eval_dataset=None, while _build_training_config sets eval_strategy="steps" with eval_steps=50. HF Trainer/SFTTrainer will raise an error the first time it attempts to evaluate with eval_strategy != "no" and no eval_dataset, failing mid-training.
🛠️ Proposed fix to retain and wire the eval split
split_dataset = dataset.train_test_split(test_size=0.1, seed=42)
self.dataset = split_dataset["train"]
+ self.eval_dataset = split_dataset["test"]
logger.info("Loaded %s training examples and %s validation examples.", len(self.dataset), len(split_dataset["test"]))
return self.dataset return SFTTrainer(
model=model,
train_dataset=self.dataset,
- eval_dataset=None,
+ eval_dataset=self.eval_dataset,
tokenizer=tokenizer,
args=self._build_training_config(),
peft_config=self._build_lora_config(),
)Also applies to: 186-193
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/train.py` around lines 71 - 74, The validation split produced in
_load_dataset is currently only logged and then discarded, while _build_trainer
still passes eval_dataset=None even though _build_training_config sets
eval_strategy to steps. Update _load_dataset to retain or return both the train
and test splits, then wire the saved validation split into _build_trainer as the
eval_dataset so Trainer/SFTTrainer can evaluate when eval_steps is reached.
| return SFTTrainer( | ||
| model=model, | ||
| train_dataset=self.dataset, | ||
| eval_dataset=None, | ||
| tokenizer=tokenizer, | ||
| args=self._build_training_config(), | ||
| peft_config=self._build_lora_config(), | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
What is the pinned version of trl in this repository's requirements.txt or pyproject.toml?
💡 Result:
The repository referred to is huggingface/trl [1][2]. In the pyproject.toml file of the huggingface/trl repository, trl itself is the project being defined, not a dependency [1]. Dependencies are listed under the [project] table [1][3]. As of the current date (July 2, 2026), trl is not listed as a pinned dependency of itself in its own pyproject.toml file; rather, it lists other libraries such as accelerate, datasets, packaging, and transformers as dependencies [3].
Citations:
- 1: https://github.com/huggingface/trl/blob/main/pyproject.toml
- 2: 📌 Pin vLLM version huggingface/trl#4122
- 3: huggingface/trl@e9aaa61
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '== pyproject / requirements mentions ==\n'
git ls-files 'pyproject.toml' 'requirements*.txt' 'uv.lock' 'poetry.lock' 'Pipfile*' | sed 's/^/- /'
printf '\n== search for trl pin/version ==\n'
rg -n --hidden --glob '!*node_modules*' --glob '!*dist*' --glob '!*build*' '\btrl\b' pyproject.toml requirements*.txt uv.lock poetry.lock Pipfile* 2>/dev/null || true
printf '\n== inspect src/train.py around the cited lines ==\n'
sed -n '160,220p' src/train.pyRepository: Chanakya-22/mini_llm_forge
Length of output: 2617
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '== requirements.txt around trl ==\n'
sed -n '1,40p' requirements.txt
printf '\n== search for SFTTrainer constructor usage/docs mentions ==\n'
rg -n '\b(processing_class|tokenizer)\b' src requirements.txt README.md . 2>/dev/null || trueRepository: Chanakya-22/mini_llm_forge
Length of output: 7931
🌐 Web query:
TRL SFTTrainer tokenizer processing_class constructor signature v0.16.0 removed tokenizer
💡 Result:
In TRL v0.16.0, SFTTrainer removed the constructor argument name tokenizer and replaced it with processing_class. 1) Evidence (v0.16.0 source/docs) - The v0.16.0 SFTTrainer docs show the argument is processing_class (optional) and its type is a tokenizer/processor; it also describes that when processing_class is None, it is loaded via AutoTokenizer.from_pretrained(model.config._name_or_path).[1] - The v0.16.0 SFTTrainer implementation includes logic: “Handle the tokenizer” and, if processing_class is None, sets processing_class = AutoTokenizer.from_pretrained(...).[2] 2) Evidence (explicit rename commit) - TRL commit “Rename trainer arg tokenizer to processing_class (#2162)” updates examples/usages from tokenizer=processor.tokenizer to processing_class=processor.tokenizer.[3] 3) What to change in your code - Before (older TRL): trainer = SFTTrainer(..., tokenizer=processor.tokenizer,...) - After (TRL v0.16.0+): trainer = SFTTrainer(..., processing_class=processor.tokenizer,...) If you previously passed tokenizer to SFTTrainer in v0.16.0, you should get a TypeError (unexpected keyword argument) and should rename that keyword to processing_class instead.[3][1]
Citations:
- 1: https://huggingface.co/docs/trl/v0.16.0/sft_trainer
- 2: https://github.com/huggingface/trl/blob/v0.16.0/trl/trainer/sft_trainer.py
- 3: huggingface/trl@47d08a9
Replace tokenizer= with processing_class= in SFTTrainer. trl is unpinned here (requirements.txt just says trl), and current SFTTrainer releases no longer accept the tokenizer kwarg, so this will raise TypeError on install.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/train.py` around lines 186 - 193, The SFTTrainer construction in train.py
is still passing tokenizer=, but newer unpinned trl releases have removed that
argument. Update the SFTTrainer call in the trainer-building method to use
processing_class= with the same tokenizer/processor object instead, and keep the
rest of the arguments unchanged so it works across current releases.
|
Note Autofix is a beta feature. Expect some limitations and changes as we gather feedback and continue to improve it. Fixes Applied SuccessfullyFixed 5 file(s) based on 12 unresolved review comments. Files modified:
Commit: The changes have been pushed to the Time taken: |
Fixed 5 file(s) based on 12 unresolved review comments. Co-authored-by: CodeRabbit <noreply@coderabbit.ai>
chatml schemas along with qlora fine tuning with non-vanishing effects to the model weights and custom data changes
Summary by CodeRabbit
New Features
Bug Fixes