"""Dataset upload → parse → validate → clean → chat-format conversion. Runs entirely on CPU (spec P7); PDF/DOCX treated as untrusted input with config-driven limits (spec §4.2). Token counts are fast estimates (chars/4) unless a cached tokenizer is available.""" import csv import hashlib import io import json import pathlib from src.config_loader import get_configs class DatasetError(Exception): pass def _limits(): return get_configs().limits.get("upload", {}) def _check_size(path: pathlib.Path): mb = path.stat().st_size / 1e6 if mb > _limits().get("max_file_mb", 50): raise DatasetError(f"File is {mb:.0f} MB — limit is {_limits().get('max_file_mb')} MB.") def _extract_pdf(path): from pypdf import PdfReader # lazy (spec §12) try: reader = PdfReader(str(path)) if reader.is_encrypted: raise DatasetError("Encrypted PDFs are not accepted.") cap = _limits().get("max_extracted_chars", 5_000_000) out, total = [], 0 for page in reader.pages: t = page.extract_text() or "" total += len(t) if total > cap: raise DatasetError(f"Extracted text exceeds {cap} characters.") out.append(t) return "\n\n".join(out) except DatasetError: raise except Exception as e: # noqa: BLE001 raise DatasetError(f"PDF could not be parsed safely: {type(e).__name__}") from e def _extract_docx(path): import docx # lazy try: d = docx.Document(str(path)) text = "\n".join(p.text for p in d.paragraphs) if len(text) > _limits().get("max_extracted_chars", 5_000_000): raise DatasetError("Extracted text exceeds the configured limit.") return text except DatasetError: raise except Exception as e: # noqa: BLE001 raise DatasetError(f"DOCX could not be parsed safely: {type(e).__name__}") from e def _rows_from_structured(path: pathlib.Path): suffix = path.suffix.lower() text = path.read_text(errors="replace") if suffix == ".csv": return list(csv.DictReader(io.StringIO(text))) if suffix == ".jsonl": return [json.loads(l) for l in text.splitlines() if l.strip()] if suffix == ".json": data = json.loads(text) if isinstance(data, dict): data = data.get("data") or data.get("rows") or [data] return data raise DatasetError(f"Unsupported structured format {suffix}") FIELD_GUESSES = { "instruction": ["instruction", "question", "prompt", "input_text", "query"], "input": ["input", "context", "passage"], "output": ["output", "answer", "response", "completion", "target", "label"], } def _guess_fields(row: dict): keys = {k.lower(): k for k in row} got = {} for role, cands in FIELD_GUESSES.items(): for c in cands: if c in keys: got[role] = keys[c] break return got def _chunk_text(text, target=1500): paras = [p.strip() for p in text.split("\n\n") if len(p.strip()) > 60] chunks, buf = [], "" for p in paras: if len(buf) + len(p) > target and buf: chunks.append(buf.strip()) buf = p else: buf += "\n\n" + p if len(buf.strip()) > 200: chunks.append(buf.strip()) return chunks def prepare(file_path: str, system_prompt: str = "") -> tuple[list[dict], dict]: """Returns (records, summary). records = [{"messages": [...]}, ...]""" path = pathlib.Path(file_path) _check_size(path) suffix = path.suffix.lower() lim = _limits() if suffix in (".csv", ".json", ".jsonl"): rows = _rows_from_structured(path) if not rows: raise DatasetError("No rows found in the file.") fields = _guess_fields(rows[0]) if "output" not in fields or ("instruction" not in fields and "input" not in fields): raise DatasetError( f"Could not identify instruction/output columns. Found: {list(rows[0].keys())}. " f"Rename columns to one of {FIELD_GUESSES['instruction']} + {FIELD_GUESSES['output']}." ) records, has_refs = [], True for r in rows: user = str(r.get(fields.get("instruction", ""), "")).strip() ctx = str(r.get(fields.get("input", ""), "")).strip() if "input" in fields else "" out = str(r.get(fields["output"], "")).strip() if not (user or ctx) or not out: continue content = f"{user}\n\n{ctx}".strip() msgs = ([{"role": "system", "content": system_prompt}] if system_prompt else []) msgs += [{"role": "user", "content": content}, {"role": "assistant", "content": out}] records.append({"messages": msgs}) elif suffix in (".txt", ".pdf", ".docx"): text = (path.read_text(errors="replace") if suffix == ".txt" else _extract_pdf(path) if suffix == ".pdf" else _extract_docx(path)) chunks = _chunk_text(text) if not chunks: raise DatasetError("No usable text extracted.") records = [{"messages": ( [{"role": "system", "content": system_prompt}] if system_prompt else []) + [ {"role": "user", "content": "Continue writing in the style and subject of this document excerpt:\n\n" + c[: len(c) // 2]}, {"role": "assistant", "content": c[len(c) // 2:]}, ]} for c in chunks] has_refs = False else: raise DatasetError(f"Unsupported file type {suffix}. Accepted: CSV, JSON, JSONL, TXT, PDF, DOCX.") # clean: dedupe + empty filter seen, cleaned, dupes = set(), [], 0 for r in records: key = hashlib.sha1(json.dumps(r, sort_keys=True).encode()).hexdigest() if key in seen: dupes += 1 continue seen.add(key) cleaned.append(r) if len(cleaned) > lim.get("max_samples", 100_000): raise DatasetError(f"{len(cleaned)} samples exceed the limit {lim.get('max_samples')}.") lengths = [sum(len(m["content"]) for m in r["messages"]) for r in cleaned] est_tokens = int(sum(lengths) / 4) if est_tokens > lim.get("max_total_tokens", 20_000_000): raise DatasetError(f"Estimated {est_tokens} tokens exceed the limit.") summary = { "samples": len(cleaned), "duplicates_removed": dupes, "est_tokens": est_tokens, "avg_tokens_per_sample": round(est_tokens / max(len(cleaned), 1), 1), "avg_chars": round(sum(lengths) / max(len(lengths), 1), 1), "has_reference_answers": has_refs, "source_file": path.name, "fingerprint": hashlib.sha256(json.dumps(cleaned[:200], sort_keys=True).encode()).hexdigest()[:16], } return cleaned, summary def save_jsonl(records: list[dict], path: pathlib.Path): path.parent.mkdir(parents=True, exist_ok=True) with open(path, "w") as f: for r in records: f.write(json.dumps(r, ensure_ascii=False) + "\n") def load_jsonl(path: pathlib.Path) -> list[dict]: return [json.loads(l) for l in path.read_text().splitlines() if l.strip()]