Spaces:
Running
Running
File size: 4,841 Bytes
c1de90b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | from __future__ import annotations
from dataclasses import dataclass
import json
from pathlib import Path
import random
from typing import Any, Iterable
@dataclass(frozen=True)
class CodeRecord:
instruction: str
input: str
output: str
line_number: int
def _as_string(value: Any, field_name: str, line_number: int) -> str:
if value is None:
return ""
if not isinstance(value, str):
raise ValueError(f"Line {line_number}: {field_name} must be a string")
return value.strip()
def load_jsonl_records(path: str | Path, *, max_records: int | None = None) -> list[CodeRecord]:
data_path = Path(path)
if not data_path.exists():
raise FileNotFoundError(f"Dataset not found: {data_path}")
records: list[CodeRecord] = []
with data_path.open("r", encoding="utf-8") as handle:
for line_number, line in enumerate(handle, start=1):
raw = line.strip()
if not raw:
continue
try:
item = json.loads(raw)
except json.JSONDecodeError as exc:
raise ValueError(f"Line {line_number}: invalid JSON: {exc}") from exc
if not isinstance(item, dict):
raise ValueError(f"Line {line_number}: record must be a JSON object")
instruction = _as_string(item.get("instruction"), "instruction", line_number)
input_text = _as_string(item.get("input", ""), "input", line_number)
output = _as_string(item.get("output"), "output", line_number)
if not instruction:
raise ValueError(f"Line {line_number}: instruction is required")
if not output:
raise ValueError(f"Line {line_number}: output is required")
records.append(
CodeRecord(
instruction=instruction,
input=input_text,
output=output,
line_number=line_number,
)
)
if max_records is not None and len(records) >= max_records:
break
if not records:
raise ValueError(f"No records found in {data_path}")
return records
def split_records(records: list[CodeRecord], validation_size: float, seed: int) -> tuple[list[CodeRecord], list[CodeRecord]]:
if validation_size <= 0:
return records, []
if validation_size >= 1:
raise ValueError("validation_size must be less than 1")
shuffled = list(records)
random.Random(seed).shuffle(shuffled)
validation_count = max(1, int(len(shuffled) * validation_size))
validation = shuffled[:validation_count]
training = shuffled[validation_count:]
if not training:
raise ValueError("validation_size leaves no training records")
return training, validation
def encode_training_record(record: CodeRecord, tokenizer: Any, max_length: int) -> dict[str, list[int]]:
user_content = record.instruction
if record.input:
user_content += f"\n{record.input}"
messages = [
{"role": "user", "content": user_content},
]
prompt_ids = tokenizer.apply_chat_template(
messages,
tokenize=True,
add_generation_prompt=True,
)
if not isinstance(prompt_ids, list):
prompt_ids = prompt_ids["input_ids"]
messages_full = messages + [{"role": "model", "content": record.output}]
full_ids = tokenizer.apply_chat_template(
messages_full,
tokenize=True,
add_generation_prompt=False,
)
if not isinstance(full_ids, list):
full_ids = full_ids["input_ids"]
input_ids = list(full_ids[:max_length])
attention_mask = [1] * len(input_ids)
labels = list(input_ids)
prompt_length = min(len(prompt_ids), len(labels))
labels[:prompt_length] = [-100] * prompt_length
if not input_ids:
raise ValueError("tokenized record is empty")
if all(label == -100 for label in labels):
raise ValueError(
f"max_length={max_length} is too short; the response was fully truncated"
)
return {
"input_ids": input_ids,
"attention_mask": attention_mask,
"labels": labels,
}
class TokenizedCodeDataset:
def __init__(self, records: Iterable[CodeRecord], tokenizer: Any, max_length: int):
self.examples = []
for record in records:
try:
example = encode_training_record(record, tokenizer, max_length)
self.examples.append(example)
except ValueError as exc:
print(f"Warning: Skipping record at line {record.line_number}: {exc}")
def __len__(self) -> int:
return len(self.examples)
def __getitem__(self, index: int) -> dict[str, list[int]]:
return self.examples[index]
|