tokendye 0.1.0__tar.gz

tokendye-0.1.0/PKG-INFO

@@ -0,0 +1,10 @@
+Metadata-Version: 2.4
+Name: tokendye
+Version: 0.1.0
+Summary: Add your description here
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+Requires-Dist: pydantic
+Requires-Dist: jsonl5>=0.1.1
+
+A python package

tokendye-0.1.0/README.md

@@ -0,0 +1 @@
+A python package

tokendye-0.1.0/dataset/__init__.py

@@ -0,0 +1,4 @@
+from .dye_dataset import DyeDataset
+from .functional import from_jsonl, from_jsonl5
+
+__all__ = ["DyeDataset", "from_jsonl", "from_jsonl5"]

tokendye-0.1.0/dataset/dye_dataset.py

@@ -0,0 +1,122 @@
+from typing import TYPE_CHECKING
+
+from torch.utils.data import Dataset
+
+if TYPE_CHECKING:
+    from ..dye_label import DyeLabel
+
+def _build_sequence(
+    data: dict,
+    tokenizer,
+    labels: list['DyeLabel']
+):
+    input_ids = []
+    dye_mask = []
+
+    messages = []  # TODO: 优化
+    for segment in data["segments"]:
+        dye = segment["dye"]
+        text = segment["text"]
+        if dye == "system":
+            messages.append({"role": "system", "content": text})
+        elif dye == "user":
+            messages.append({"role": "user", "content": text})
+        elif dye == "file_text":
+            messages.append({"role": "user", "content": text})
+        elif dye == "tool_callback":
+            messages.append({"role": "tool", "content": text})
+        else:
+            messages.append({"role": "user", "content": text})
+
+    # tokenize=True直接返回ids，不含generation prompt
+    full_ids: list[int] = tokenizer.apply_chat_template(
+        messages,
+        tokenize=True,
+        add_generation_prompt=False,
+    )["input_ids"]  # TODO: 留着attention_mask或许有用
+
+    # 2. 对每段content单独encode，在full_ids里定位边界，构建dye_mask
+    dye_mask = [-1] * len(full_ids)
+
+    label_map = {label.name:label.id for label in labels}
+    for segment in data["segments"]:
+        dye = segment["dye"]
+        text = segment["text"]
+        dye_id = label_map[dye]
+
+        content_ids = tokenizer.encode(text, add_special_tokens=False)
+        positions = _find_all_sublist(full_ids, content_ids)
+
+        if len(positions) == 0:
+            raise ValueError(f"content not found in full_ids: {text!r}")
+        if len(positions) > 1:
+            raise ValueError(
+                f"ambiguous match ({len(positions)} hits), "
+                f"two segments have identical content: {text!r}"
+            )
+
+        pos = positions[0]
+        for i in range(len(content_ids)):
+            dye_mask[pos + i] = dye_id
+
+    # 3. target部分：从apply_chat_template提取generation prompt
+    context_len = len(full_ids)
+
+    full_ids_with_gen: list[int] = tokenizer.apply_chat_template(
+        messages,
+        tokenize=True,
+        add_generation_prompt=True,
+    )["input_ids"]  # TODO: 留着attention_mask或许有用
+    gen_prompt_ids = full_ids_with_gen[context_len:]  # 只取新增部分
+
+    target_ids = tokenizer.encode(data["target"], add_special_tokens=False)
+    eos = [tokenizer.eos_token_id]
+
+    input_ids = full_ids + gen_prompt_ids + target_ids + eos
+    dye_mask = dye_mask + [-1] * (len(gen_prompt_ids) + len(target_ids) + 1)
+
+    target_mask = (
+        [False] * context_len
+        + [False] * len(gen_prompt_ids)
+        + [True] * len(target_ids)
+        + [True]  # eos算loss
+    )
+
+    return input_ids, dye_mask, target_mask
+
+
+class DyeDataset(Dataset):
+    def __init__(
+        self,
+        raw_data: list[dict],
+        tokenizer,
+        labels: list['DyeLabel']
+    ):
+        self.dataset: list[dict] = []
+
+        for data in raw_data:
+            input_ids, dye_mask, target_mask = _build_sequence(
+                data, tokenizer, labels
+            )
+
+            self.dataset.append(
+                {
+                    "input_ids": input_ids,
+                    "dye_mask": dye_mask,
+                    "target_mask": target_mask,
+                }
+            )
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        return self.dataset[idx]
+
+
+def _find_all_sublist(full, sub) -> list[int]:
+    """返回sub在full中所有出现位置的起始index"""
+    n, m = len(full), len(sub)
+    if m == 0:
+        return []
+    return [i for i in range(n - m + 1) if full[i : i + m] == sub]

tokendye-0.1.0/dataset/functional.py

@@ -0,0 +1,28 @@
+import json
+from typing import TYPE_CHECKING
+
+import jsonl5
+
+from . import DyeDataset
+
+if TYPE_CHECKING:
+    from os import PathLike
+
+    from ..dye_label import DyeLabel
+
+
+def from_jsonl(data_path: "PathLike | str", tokenizer, labels: list["DyeLabel"]):
+    with open(data_path) as f:
+        # raw_data = []
+        # for line in f:
+        #     raw_data.append(json.loads(line))
+        raw_data = [json.loads(line) for line in f]
+
+    return DyeDataset(raw_data, tokenizer, labels)
+
+
+def from_jsonl5(data_path: "PathLike | str", tokenizer, labels: list["DyeLabel"]):
+    with open(data_path) as f:
+        raw_data = jsonl5.load(f)
+
+    return DyeDataset(raw_data, tokenizer, labels)

tokendye-0.1.0/pyproject.toml

@@ -0,0 +1,10 @@
+[project]
+name = "tokendye"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "pydantic",
+    "jsonl5>=0.1.1",
+]

tokendye-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

tokendye-0.1.0/tokendye.egg-info/PKG-INFO

@@ -0,0 +1,10 @@
+Metadata-Version: 2.4
+Name: tokendye
+Version: 0.1.0
+Summary: Add your description here
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+Requires-Dist: pydantic
+Requires-Dist: jsonl5>=0.1.1
+
+A python package

tokendye-0.1.0/tokendye.egg-info/SOURCES.txt

@@ -0,0 +1,10 @@
+README.md
+pyproject.toml
+dataset/__init__.py
+dataset/dye_dataset.py
+dataset/functional.py
+tokendye.egg-info/PKG-INFO
+tokendye.egg-info/SOURCES.txt
+tokendye.egg-info/dependency_links.txt
+tokendye.egg-info/requires.txt
+tokendye.egg-info/top_level.txt

tokendye-0.1.0/tokendye.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

tokendye-0.1.0/tokendye.egg-info/requires.txt

@@ -0,0 +1,2 @@
+pydantic
+jsonl5>=0.1.1

tokendye-0.1.0/tokendye.egg-info/top_level.txt

@@ -0,0 +1 @@
+dataset