Replace fake train loop with TRL SFTTrainer, HF column mapping with fiction preset, safetensors to GGUF conversion, and ollama create using ollama_base plus ADAPTER. Co-authored-by: Cursor <cursoragent@cursor.com>
95 lines
3.4 KiB
Python
95 lines
3.4 KiB
Python
"""Map Hugging Face dataset rows to chat messages for SFT."""
|
|
from __future__ import annotations
|
|
|
|
|
|
def fiction_tags_text_user(row: dict, title_col: str = "title", tags_col: str = "tags") -> str:
|
|
parts = []
|
|
title = (row.get(title_col) or "").strip()
|
|
tags = (row.get(tags_col) or "").strip()
|
|
if title:
|
|
parts.append(f"Title: {title}")
|
|
if tags:
|
|
parts.append(f"Tags: {tags}")
|
|
return "\n".join(parts) if parts else tags or title or ""
|
|
|
|
|
|
def row_to_messages(row: dict, schema: dict | None, mapping: dict | None) -> list[dict] | None:
|
|
schema = schema or {}
|
|
mapping = mapping or {}
|
|
kind = schema.get("kind") or mapping.get("kind") or mapping.get("preset")
|
|
|
|
if kind == "messages" and row.get("messages"):
|
|
return _normalize_messages(row["messages"])
|
|
|
|
if kind == "conversations" and row.get("conversations"):
|
|
out = []
|
|
for c in row["conversations"]:
|
|
if not isinstance(c, dict):
|
|
continue
|
|
frm = c.get("from") or ""
|
|
val = c.get("value") or ""
|
|
role = "assistant" if frm in ("gpt", "assistant", "chatgpt") else "user"
|
|
if frm in ("human", "user"):
|
|
role = "user"
|
|
if val:
|
|
out.append({"role": role, "content": str(val)})
|
|
return out or None
|
|
|
|
if kind == "alpaca":
|
|
instr = (row.get("instruction") or "").strip()
|
|
inp = (row.get("input") or "").strip()
|
|
output = (row.get("output") or "").strip()
|
|
user = instr if not inp else f"{instr}\n{inp}"
|
|
if user and output:
|
|
return [{"role": "user", "content": user}, {"role": "assistant", "content": output}]
|
|
return None
|
|
|
|
if kind == "prompt_response":
|
|
resp_col = schema.get("response_col") or "response"
|
|
prompt = row.get("prompt") or ""
|
|
resp = row.get(resp_col) or ""
|
|
if prompt and resp:
|
|
return [{"role": "user", "content": str(prompt)}, {"role": "assistant", "content": str(resp)}]
|
|
return None
|
|
|
|
if kind == "qa":
|
|
q = row.get("question") or ""
|
|
a = row.get("answer") or ""
|
|
if q and a:
|
|
return [{"role": "user", "content": str(q)}, {"role": "assistant", "content": str(a)}]
|
|
return None
|
|
|
|
if kind in ("fiction_tags_text", "preset_fiction_tags_text"):
|
|
text = (row.get("text") or "").strip()
|
|
user = fiction_tags_text_user(row)
|
|
if user and text:
|
|
return [{"role": "user", "content": user}, {"role": "assistant", "content": text}]
|
|
return None
|
|
|
|
if kind == "custom" or mapping.get("user_col"):
|
|
user_col = mapping.get("user_col")
|
|
asst_col = mapping.get("assistant_col")
|
|
if user_col and asst_col:
|
|
u = row.get(user_col) or ""
|
|
a = row.get(asst_col) or ""
|
|
if u and a:
|
|
return [{"role": "user", "content": str(u)}, {"role": "assistant", "content": str(a)}]
|
|
return None
|
|
|
|
if row.get("messages"):
|
|
return _normalize_messages(row["messages"])
|
|
|
|
return None
|
|
|
|
|
|
def _normalize_messages(msgs) -> list[dict] | None:
|
|
out = []
|
|
for m in msgs:
|
|
if not isinstance(m, dict):
|
|
continue
|
|
role = m.get("role") or "user"
|
|
content = m.get("content") or m.get("text") or ""
|
|
if content:
|
|
out.append({"role": role, "content": str(content)})
|
|
return out or None
|