Spaces:
Running
Running
| """DataInventor engine: one sentence in, a training dataset out, on any OpenAI-compatible endpoint. | |
| Pipeline: blueprint (taxonomy + personas + question shapes) -> sampled combinations -> | |
| question -> answer (or two answers + a ranking for preference pairs) -> filters -> near-duplicate rejection. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import random | |
| import re | |
| import threading | |
| import time | |
| from concurrent.futures import ThreadPoolExecutor | |
| from dataclasses import dataclass, field | |
| from typing import Optional | |
| from datainventor.llm import HF_BASE, HF_MODEL, Endpoint, EndpointError # noqa: F401 (main imports them from here) | |
| # ---------------------------------------------------------------- blueprint | |
| def _s(x) -> str: | |
| """Models sometimes return {"persona": "..."} instead of a bare string.""" | |
| if isinstance(x, dict): | |
| for k in ("name", "persona", "role", "format", "shape", "topic", "description", "title"): | |
| if isinstance(x.get(k), str): | |
| return x[k].strip() | |
| return " ".join(str(v) for v in x.values() if isinstance(v, str)).strip() | |
| return str(x).strip() | |
| def normalize_blueprint(bp: dict) -> dict: | |
| out = {"topics": [], "personas": [_s(p) for p in bp.get("personas", []) if _s(p)], | |
| "formats": [_s(f) for f in bp.get("formats", []) if _s(f)], | |
| "difficulty": [_s(d) for d in bp.get("difficulty", []) if _s(d)] or ["basic", "intermediate", "advanced"], | |
| "assistant_role": _s(bp.get("assistant_role") or "an excellent, accurate AI assistant")[:300]} | |
| for a in bp.get("topics", []): | |
| subs = [] | |
| for sub in a.get("subtopics", []): | |
| leaves = [_s(l) for l in sub.get("leaf_topics", []) if _s(l)] | |
| if leaves: | |
| subs.append({"subtopic": _s(sub.get("subtopic", "")), "leaf_topics": leaves}) | |
| if subs: | |
| out["topics"].append({"area": _s(a.get("area", "")), "subtopics": subs}) | |
| if len(out["topics"]) < 1 or len(out["personas"]) < 2 or len(out["formats"]) < 2: | |
| raise ValueError("the blueprint is missing topics, personas or formats") | |
| return out | |
| def _json_from(text: str): | |
| m = re.search(r"\{.*\}", text, re.S) | |
| if not m: | |
| raise ValueError("no JSON object in the reply") | |
| return json.loads(m.group(0)) | |
| def make_blueprint(ep: Endpoint, request: str, domains: list[str], subdomains: list[str], languages: list[str], | |
| constraints: str, rows: int) -> dict: | |
| scope = "" | |
| if domains: | |
| scope += f"Domains: {', '.join(domains)}.\n" | |
| if subdomains: | |
| scope += f"Subdomains: {', '.join(subdomains)}.\n" | |
| if constraints.strip(): | |
| scope += f"Output requirements: {constraints.strip()}\n" | |
| areas = 6 if rows <= 500 else 8 | |
| prompt = f"""You are planning a training dataset. Dataset request: "{request}" | |
| {scope} | |
| Return ONLY JSON with these keys: | |
| "topics": a list of {areas} objects {{"area": str, "subtopics": [4 objects {{"subtopic": str, "leaf_topics": [4 specific, concrete strings]}}]}}. Cover the whole request evenly. | |
| "personas": 10 different kinds of people who would write the user side of this dataset (for a customer-support dataset, 10 different kinds of customer; for a coding dataset, 10 kinds of developer). | |
| "assistant_role": one sentence describing who answers in this dataset and in what voice (for example "a friendly support agent at a small online bookshop who gives concrete next steps"). | |
| "formats": 6 different shapes the instruction could take (for example a short direct ask, a detailed brief, a request with a concrete example, a request with constraints, a follow-up, a comparison). | |
| "difficulty": ["basic","intermediate","advanced"]. | |
| No commentary, no code fences.""" | |
| last = None | |
| for _ in range(4): | |
| try: | |
| return normalize_blueprint(_json_from(ep.chat(prompt, max_tokens=5000, temperature=0.7))) | |
| except EndpointError: | |
| raise | |
| except Exception as e: | |
| last = e | |
| raise EndpointError(f"The model did not return a usable blueprint: {last}") | |
| # ---------------------------------------------------------------- generation | |
| LEAK = re.compile(r"the user (asks|wants)|write one|return only|^okay|^sure,|^here is|^here's|^as an ai", re.I) | |
| END = re.compile(r"[.?!)\"”’'।:。！？]$") | |
| MULTI = re.compile(r"\b(generate|draft|write|compile|create|give me|produce)\s+(\d+|several|multiple|a set of|a list of)\b", re.I) | |
| PLACEHOLDER = re.compile(r"\[[A-Za-z'’ ]{0,30}\b(name|number|email|address|phone|date|id)\]", re.I) | |
| CALLBACK = re.compile(r"\b(like|as) you (said|mentioned|suggested|told me)\b|\b(last time|earlier) (you|we)\b|\bthanks again\b", re.I) | |
| TABLE = re.compile(r"^\s*\|.*\|\s*$\n^\s*\|?\s*:?-{3,}", re.M) | |
| VERDICT = re.compile(r"\W*\(?([AB])\b", re.I) | |
| LENGTHS = ["one blunt sentence", "two sentences with one specific detail", | |
| "three or four sentences with concrete specifics", "a short paragraph that sounds a little rushed"] | |
| def degenerate(t: str) -> bool: | |
| w = t.lower().split() | |
| return bool(re.search(r"(.)\1{7,}", t)) or (len(w) >= 8 and len(set(w)) / len(w) < 0.5) | |
| def grams(t: str, n: int = 3) -> set: | |
| w = t.lower().split() | |
| return set(zip(*[w[i:] for i in range(n)])) if len(w) >= n else {tuple(w)} | |
| class Job: | |
| id: str | |
| config: dict | |
| state: str = "queued" # queued | planning | generating | done | failed | stopped | |
| message: str = "" | |
| blueprint: Optional[dict] = None | |
| rows: list = field(default_factory=list) | |
| target: int = 0 | |
| attempts: int = 0 | |
| rejected: dict = field(default_factory=dict) | |
| started: float = field(default_factory=time.time) | |
| finished: float = 0.0 | |
| stop: threading.Event = field(default_factory=threading.Event) | |
| lock: threading.Lock = field(default_factory=threading.Lock) | |
| def reject(self, why: str): | |
| with self.lock: | |
| self.rejected[why] = self.rejected.get(why, 0) + 1 | |
| def run_job(job: Job, ep: Endpoint): | |
| cfg = job.config | |
| rows_n = int(cfg["rows"]) | |
| job.target = rows_n | |
| kind = cfg.get("format", "instruction") | |
| request = cfg["request"] | |
| languages = cfg.get("languages") or ["English"] | |
| constraints = cfg.get("constraints", "") | |
| try: | |
| job.state, job.message = "planning", "Drafting the blueprint" | |
| bp = normalize_blueprint(cfg["blueprint"]) if cfg.get("blueprint") else make_blueprint( | |
| ep, request, cfg.get("domains", []), cfg.get("subdomains", []), languages, constraints, rows_n) | |
| job.blueprint = bp | |
| leaves = [(a["area"], s["subtopic"], l) for a in bp["topics"] for s in a["subtopics"] for l in s["leaf_topics"]] | |
| rng = random.Random(cfg.get("seed", 7)) | |
| combos, seen = [], set() | |
| want = int(rows_n * 1.8) + 20 | |
| guard = 0 | |
| while len(combos) < want and guard < want * 30: | |
| guard += 1 | |
| c = (rng.choice(leaves), rng.choice(bp["personas"]), rng.choice(bp["formats"]), rng.choice(bp["difficulty"]), | |
| rng.choice(LENGTHS), rng.choice(languages)) | |
| if c not in seen: | |
| seen.add(c) | |
| combos.append(c) | |
| scope = f'The dataset being built: "{request}"' | |
| if constraints.strip(): | |
| scope += f"\nEvery example must satisfy: {constraints.strip()}" | |
| def ok_q(t): | |
| t = t.strip().strip('"') | |
| return (15 <= len(t) <= 1500 and END.search(t) and not degenerate(t) and not LEAK.search(t[:150]) | |
| and not MULTI.search(t[:120]) and not PLACEHOLDER.search(t) and not CALLBACK.search(t)) | |
| def ok_a(t): | |
| return (len(t.split()) >= 12 and not degenerate(t[:600]) and not LEAK.search(t[:100]) | |
| and not TABLE.search(t) and not PLACEHOLDER.search(t)) | |
| kept_grams: list[set] = [] | |
| job.state, job.message = "generating", "Writing examples" | |
| def work(idx_combo): | |
| idx, ((area, sub, leaf), persona, fmt, diff, length, lang) = idx_combo | |
| if job.stop.is_set(): | |
| return | |
| with job.lock: | |
| if len(job.rows) >= rows_n: | |
| return | |
| job.attempts += 1 | |
| try: | |
| role = bp.get("assistant_role", "an excellent AI assistant") | |
| q = ep.chat(f"""{scope} | |
| Write ONE realistic message that {persona} would send to {role}, for this dataset, about: {leaf} ({area}, {sub}). | |
| Shape: {fmt}. Difficulty: {diff}. Length: {length}. Language: {lang}. | |
| It should sound like a real person typing, specific and a little imperfect, not a textbook prompt. It must be ONE question or ONE task whose answer fits in a single reply: never ask for a list of many items, a set of Q&A pairs, or a dataset. Do not name the shape or the persona. | |
| It is the first message in the conversation: do not refer to earlier chats or to anything said before. Do not use placeholders such as [Name] or [order number]. | |
| Return only the message.""", max_tokens=700, temperature=1.0, validate=ok_q).strip().strip('"') | |
| g = grams(q) | |
| with job.lock: | |
| if any(len(g & k) / max(1, len(g | k)) > 0.5 for k in kept_grams): | |
| job.reject("near-duplicate question") | |
| return | |
| answer_prompt = f"""You are {role}. Reply to the message below. The reply is a training example for this dataset: "{request}". | |
| Language: {lang}. Be accurate, specific and complete. Start directly with the content, with no preamble about who you are. Do not invent statistics or facts you cannot know. | |
| Write plain prose: short paragraphs, or a simple numbered list for steps. No tables, headings or emoji. | |
| Never invent order numbers, dates, prices, phone numbers, emails, URLs, product or brand names, menu paths or policies. If a detail like that is needed, ask for it or speak generally. Do not use placeholders such as [Your Name]. | |
| For physical techniques and safety, give only the standard, widely taught practice, and make sure every step agrees with the others. | |
| {('Requirements: ' + constraints.strip()) if constraints.strip() else ''} | |
| Message: {q}""" | |
| a1 = ep.chat(answer_prompt, max_tokens=2500, temperature=0.7, validate=ok_a) | |
| row = {"area": area, "subtopic": sub, "topic": leaf, "persona": persona, "shape": fmt, | |
| "difficulty": diff, "language": lang, "prompt": q} | |
| if kind == "preference": | |
| a2 = ep.chat(answer_prompt + "\n\nWrite a different answer from the obvious one, in a different structure and with different emphasis.", | |
| max_tokens=2500, temperature=1.0, validate=ok_a) | |
| def judge(first: str, second: str) -> str: | |
| v = ep.chat(f"""Which answer is better for the message? Judge accuracy first: an answer that invents details (order numbers, dates, contact details, policies) or gives unsafe advice is worse. Then completeness and clarity. | |
| Message: {q} | |
| Answer A: | |
| {first[:3000]} | |
| Answer B: | |
| {second[:3000]} | |
| Reply with only the single letter A or B.""", max_tokens=16, temperature=0, allow_cut=True, | |
| validate=lambda s: bool(VERDICT.match(s))) | |
| return VERDICT.match(v).group(1).upper() | |
| # judge both orders so the first-shown answer gets no edge; keep the row only if they agree | |
| v1, v2 = judge(a1, a2), judge(a2, a1) | |
| if (v1 == "A") != (v2 == "B"): | |
| job.reject("judge disagreed") | |
| return | |
| row["chosen"], row["rejected"] = (a1, a2) if v1 == "A" else (a2, a1) | |
| else: | |
| row["completion"] = a1 | |
| with job.lock: | |
| if len(job.rows) >= rows_n: | |
| return | |
| if any(len(g & k) / max(1, len(g | k)) > 0.5 for k in kept_grams): | |
| job.reject("near-duplicate question") | |
| return | |
| row["id"] = len(job.rows) + 1 | |
| job.rows.append(row) | |
| kept_grams.append(g) | |
| job.message = f"{len(job.rows)} of {rows_n} written" | |
| if len(job.rows) >= rows_n: | |
| job.state, job.message, job.finished = "done", f"{len(job.rows)} rows ready", time.time() | |
| job.stop.set() | |
| except EndpointError as e: | |
| if "refused the key" in str(e) or "not found" in str(e): | |
| job.state, job.message = "failed", str(e) | |
| job.stop.set() | |
| else: | |
| job.reject("model call failed") | |
| except Exception: | |
| job.reject("model call failed") | |
| with ThreadPoolExecutor(max_workers=max(1, min(ep.workers, 16))) as pool: | |
| list(pool.map(work, enumerate(combos))) | |
| if job.state in ("failed", "done"): | |
| return | |
| if job.stop.is_set(): | |
| job.state, job.message = "stopped", f"Stopped with {len(job.rows)} rows" | |
| elif len(job.rows) < rows_n: | |
| job.state = "done" | |
| job.message = f"Finished with {len(job.rows)} of {rows_n} rows. The model rejected too many attempts to fill the rest." | |
| else: | |
| job.state, job.message = "done", f"{len(job.rows)} rows ready" | |
| except EndpointError as e: | |
| job.state, job.message = "failed", str(e) | |
| except Exception as e: # never leave a job hanging | |
| job.state, job.message = "failed", f"Unexpected error: {e}" | |
| finally: | |
| job.finished = time.time() | |
| # ---------------------------------------------------------------- stats | |
| def stats(rows: list) -> dict: | |
| if not rows: | |
| return {} | |
| text = "\n".join(r["prompt"] for r in rows).lower().split() | |
| tri = list(zip(text, text[1:], text[2:])) | |
| area_counts: dict = {} | |
| for r in rows: | |
| area_counts[r["area"]] = area_counts.get(r["area"], 0) + 1 | |
| comp = [len((r.get("completion") or r.get("chosen") or "").split()) for r in rows] | |
| return {"rows": len(rows), "distinct3_prompts": round(len(set(tri)) / max(1, len(tri)), 3), | |
| "avg_prompt_words": round(sum(len(r["prompt"].split()) for r in rows) / len(rows), 1), | |
| "avg_answer_words": round(sum(comp) / len(rows), 1), "areas": area_counts} | |