rpeft 0.1.0

data/python/rpeft_backend/__init__.py

@@ -0,0 +1,90 @@
+"""Private Python backend for the RPEFT Ruby gem.
+
+Imported through PyCall; not a public API. Every facade function takes
+a JSON string and returns a JSON string, so only plain strings cross
+the PyCall boundary.
+
+This module must stay importable without torch installed: all heavy
+imports live inside function bodies, and diagnostics() reports what is
+missing instead of failing.
+"""
+
+import json
+import math
+import os
+import sys
+
+__version__ = "0.1.0"
+
+
+def _dumps(value):
+    # json.dumps writes NaN and Infinity, which are not JSON and which the Ruby side cannot parse.
+    # A diverged run reports such losses, so they cross as null rather than failing the whole call.
+    return json.dumps(_finite(value), allow_nan=False)
+
+
+def _finite(value):
+    if isinstance(value, float) and not math.isfinite(value):
+        return None
+    if isinstance(value, dict):
+        return {key: _finite(item) for key, item in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [_finite(item) for item in value]
+    return value
+
+
+def _repair_sys_executable():
+    # PyCall embeds Python inside the Ruby process, so sys.executable
+    # points at the Ruby binary. multiprocessing (used by datasets for
+    # its resource tracker) would then spawn Ruby instead of Python.
+    # Repair it from the PYTHON environment variable that PyCall honors.
+    python = os.environ.get("PYTHON")
+    if not python or not os.path.isfile(python):
+        return
+    if "python" in os.path.basename(sys.executable):
+        return
+
+    import multiprocessing
+
+    sys.executable = python
+    multiprocessing.set_executable(python)
+
+
+_repair_sys_executable()
+
+
+def train_lora(config_json):
+    from . import _training, _versions
+
+    _versions.require_minimums()
+
+    return _dumps(_training.train_lora(json.loads(config_json)))
+
+
+def merge_lora(config_json):
+    from . import _merge, _versions
+
+    _versions.require_minimums()
+
+    return _dumps(_merge.merge_lora(json.loads(config_json)))
+
+
+def generate(config_json):
+    from . import _generate, _versions
+
+    _versions.require_minimums()
+
+    return _dumps(_generate.generate(json.loads(config_json)))
+
+
+def detect_device(config_json="{}"):
+    from . import _device
+
+    config = json.loads(config_json)
+    return _dumps(_device.resolve_device(config.get("device")))
+
+
+def diagnostics():
+    from . import _diagnostics
+
+    return _dumps(_diagnostics.diagnostics())

data/python/rpeft_backend/_chat.py

@@ -0,0 +1,79 @@
+def render(tokenizer, messages, add_generation_prompt=False):
+    """Renders a chat transcript into the text the model was trained to see.
+
+    add_generation_prompt marks where the assistant's turn begins. It belongs to
+    inference only: training renders the transcript as written.
+    """
+    _check_roles(messages)
+    messages = _without_control_tokens(tokenizer, messages)
+    if getattr(tokenizer, "chat_template", None):
+        return tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=add_generation_prompt
+        )
+
+    # Tokenizers without a chat template (small test models, some base models) still need
+    # a deterministic rendering.
+    rendered = "\n".join(
+        "%s: %s" % (message["role"], message["content"]) for message in messages
+    )
+    if add_generation_prompt:
+        rendered += "\nassistant:"
+
+    return rendered
+
+
+def writes_special_tokens(tokenizer):
+    """Whether render's output already holds the special tokens, BOS included, that the model expects.
+
+    A chat template writes them itself, so tokenizing its output must not add them again; the plain
+    fallback rendering writes none.
+    """
+    return bool(getattr(tokenizer, "chat_template", None))
+
+
+# Kept in step with RPEFT::Messages::ROLES. The Ruby side checks only the first line of a dataset file,
+# so the rest of the file is checked here.
+ROLES = ("system", "user", "assistant")
+
+
+def _check_roles(messages):
+    for message in messages:
+        if message.get("role") not in ROLES:
+            raise ValueError(
+                "message role must be one of %s, got %r" % (", ".join(ROLES), message.get("role"))
+            )
+
+
+def _without_control_tokens(tokenizer, messages):
+    # The rendered text is tokenized as a whole, so a control token such as <|im_end|> written inside
+    # a message would end the turn there and the rest would be read as a turn of its own. Message text
+    # is data, so such tokens are removed from it; only the chat template may produce them.
+    tokens = _control_tokens(tokenizer)
+    if not tokens:
+        return messages
+
+    return [
+        dict(message, content=_strip(message["content"], tokens))
+        if isinstance(message.get("content"), str)
+        else message
+        for message in messages
+    ]
+
+
+def _control_tokens(tokenizer):
+    tokens = set(getattr(tokenizer, "all_special_tokens", None) or [])
+    for token in (getattr(tokenizer, "added_tokens_decoder", None) or {}).values():
+        if getattr(token, "special", False):
+            tokens.add(token.content)
+    return sorted(tokens, key=len, reverse=True)
+
+
+def _strip(content, tokens):
+    # Repeated until stable, so removing one token cannot join its neighbors into another.
+    while True:
+        stripped = content
+        for token in tokens:
+            stripped = stripped.replace(token, "")
+        if stripped == content:
+            return stripped
+        content = stripped

data/python/rpeft_backend/_data.py

@@ -0,0 +1,59 @@
+import json
+
+from . import _chat
+from ._errors import BackendError, stage
+
+
+def load_tokenized_dataset(path, tokenizer, text_field, max_length, dataset_format="text"):
+    with stage("dataset"):
+        from datasets import Dataset
+
+        # Built in memory rather than with load_dataset("json", ...), which writes an Arrow copy of every
+        # file into the Hugging Face cache. Each run reads a fresh file, so that copy would never be reused.
+        dataset = Dataset.from_list(_read_jsonl(path))
+        field = "messages" if dataset_format == "messages" else text_field
+        if field not in dataset.column_names:
+            raise BackendError(
+                "dataset",
+                "field '%s' not found in dataset (available: %s)"
+                % (field, ", ".join(dataset.column_names)),
+            )
+
+        add_special_tokens = not (dataset_format == "messages" and _chat.writes_special_tokens(tokenizer))
+        # A chat template closes every turn with its own end marker. Plain text has none, so without an EOS
+        # the model never sees where an answer ends and keeps generating until max_new_tokens.
+        eos = tokenizer.eos_token_id if dataset_format != "messages" else None
+        content_length = max_length - 1 if eos is not None else max_length
+
+        def tokenize(batch):
+            if dataset_format == "messages":
+                texts = [_chat.render(tokenizer, messages) for messages in batch[field]]
+            else:
+                texts = batch[field]
+
+            # No padding here: the collator pads each batch dynamically.
+            encoded = tokenizer(
+                texts, truncation=True, max_length=content_length, add_special_tokens=add_special_tokens
+            )
+            if eos is not None:
+                for input_ids, attention_mask in zip(encoded["input_ids"], encoded["attention_mask"]):
+                    input_ids.append(eos)
+                    attention_mask.append(1)
+            return encoded
+
+        return dataset.map(
+            tokenize, batched=True, remove_columns=dataset.column_names, keep_in_memory=True
+        )
+
+
+def _read_jsonl(path):
+    records = []
+    with open(path, encoding="utf-8") as lines:
+        for lineno, line in enumerate(lines, 1):
+            if not line.strip():
+                continue
+            try:
+                records.append(json.loads(line))
+            except json.JSONDecodeError as error:
+                raise BackendError("dataset", "line %d is not valid JSON: %s" % (lineno, error))
+    return records

data/python/rpeft_backend/_device.py

@@ -0,0 +1,57 @@
+from ._errors import BackendError, stage
+
+_KNOWN_DEVICES = ("cuda", "mps", "cpu")
+
+
+def resolve_device(override=None):
+    with stage("device"):
+        import torch
+
+        if override in (None, "", "auto"):
+            device = _detect(torch)
+            reason = "cpu fallback" if device == "cpu" else "%s available" % device
+        else:
+            device = _validate(torch, override)
+            reason = "explicit override"
+
+        return {
+            "device": device,
+            "dtype": _dtype_for(torch, device),
+            "reason": reason,
+            "device_name": torch.cuda.get_device_name(0) if device == "cuda" else None,
+        }
+
+
+def _detect(torch):
+    if torch.cuda.is_available():
+        return "cuda"
+    if torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
+
+
+def _validate(torch, requested):
+    if requested not in _KNOWN_DEVICES:
+        raise BackendError("device", "unknown device: %s" % requested)
+    if requested == "cuda" and not torch.cuda.is_available():
+        raise BackendError("device", "cuda requested but not available")
+    if requested == "mps" and not torch.backends.mps.is_available():
+        raise BackendError("device", "mps requested but not available")
+    return requested
+
+
+def _dtype_for(torch, device):
+    if device == "cuda" and torch.cuda.is_bf16_supported():
+        return "bfloat16"
+    # fp16/bf16 on MPS are numerically fragile; v0.1 sticks to float32
+    # everywhere except bf16-capable CUDA.
+    return "float32"
+
+
+def empty_cache(device):
+    import torch
+
+    if device == "mps" and hasattr(torch, "mps"):
+        torch.mps.empty_cache()
+    elif device == "cuda":
+        torch.cuda.empty_cache()

data/python/rpeft_backend/_diagnostics.py

@@ -0,0 +1,77 @@
+import importlib
+import importlib.util
+import platform
+import sys
+
+LIBRARIES = ("torch", "transformers", "peft", "datasets", "accelerate")
+
+
+def diagnostics():
+    return {
+        "python": {
+            "executable": sys.executable,
+            "version": platform.python_version(),
+        },
+        "backend_version": _backend_version(),
+        "libraries": {name: _library_info(name) for name in LIBRARIES},
+        "device": _device_info(),
+    }
+
+
+def _backend_version():
+    from . import __version__
+
+    return __version__
+
+
+def _library_info(name):
+    from . import _versions
+
+    minimum = _versions.MINIMUMS.get(name)
+    if importlib.util.find_spec(name) is None:
+        return {"installed": False, "version": None, "minimum": minimum, "outdated": False}
+    try:
+        module = importlib.import_module(name)
+        version = getattr(module, "__version__", None)
+        return {
+            "installed": True,
+            "version": version,
+            "minimum": minimum,
+            "outdated": _versions.is_outdated(name, version),
+        }
+    except Exception as error:
+        return {
+            "installed": False,
+            "version": None,
+            "minimum": minimum,
+            "outdated": False,
+            "error": str(error),
+        }
+
+
+def _device_info():
+    try:
+        import torch
+    except Exception:
+        return None
+
+    cuda_available = torch.cuda.is_available()
+    mps_available = torch.backends.mps.is_available()
+    if cuda_available:
+        recommended = "cuda"
+    elif mps_available:
+        recommended = "mps"
+    else:
+        recommended = "cpu"
+
+    return {
+        "cuda": {
+            "available": cuda_available,
+            "device_name": torch.cuda.get_device_name(0) if cuda_available else None,
+        },
+        "mps": {
+            "available": mps_available,
+            "built": torch.backends.mps.is_built(),
+        },
+        "recommended": recommended,
+    }

data/python/rpeft_backend/_errors.py

@@ -0,0 +1,24 @@
+import contextlib
+
+
+class BackendError(RuntimeError):
+    """Error whose message carries a machine-readable stage token.
+
+    The Ruby side parses the "[stage:...]" prefix to pick the matching
+    Ruby exception class; see lib/rpeft/backend/error_translation.rb.
+    """
+
+    def __init__(self, stage, message):
+        super().__init__("[stage:%s] %s" % (stage, message))
+        self.stage = stage
+
+
+@contextlib.contextmanager
+def stage(name):
+    """Tag any exception escaping this block with the given stage name."""
+    try:
+        yield
+    except BackendError:
+        raise
+    except Exception as error:
+        raise BackendError(name, "%s: %s" % (type(error).__name__, error)) from error

data/python/rpeft_backend/_generate.py

@@ -0,0 +1,73 @@
+import gc
+
+from . import _chat, _device, _models
+from ._errors import stage
+
+
+def generate(config):
+    resolved = _device.resolve_device(config.get("device"))
+    device = resolved["device"]
+    dtype = resolved["dtype"]
+
+    model = None
+    try:
+        tokenizer, model = _models.load_tokenizer_and_model(
+            config["model"],
+            dtype=dtype,
+            device=device,
+            trust_remote_code=config.get("trust_remote_code", False),
+        )
+
+        adapter_path = config.get("adapter_path")
+        if adapter_path:
+            with stage("model_load"):
+                from peft import PeftModel
+
+                # Applied on top of the base weights; not merged.
+                model = PeftModel.from_pretrained(model, adapter_path)
+
+        with stage("generation"):
+            import torch
+
+            if config.get("seed") is not None:
+                torch.manual_seed(config["seed"])
+
+            # Greedy decoding by default keeps the output deterministic
+            # for a given model and prompt, which demo scripts rely on;
+            # temperature or top_p switches sampling on.
+            generation_options = {
+                "max_new_tokens": config.get("max_new_tokens", 64),
+                "do_sample": False,
+                "pad_token_id": tokenizer.pad_token_id,
+            }
+            if config.get("temperature") is not None:
+                generation_options["temperature"] = float(config["temperature"])
+                generation_options["do_sample"] = True
+            if config.get("top_p") is not None:
+                generation_options["top_p"] = float(config["top_p"])
+                generation_options["do_sample"] = True
+
+            messages = config.get("messages")
+            if messages:
+                prompt = _chat.render(tokenizer, messages, add_generation_prompt=True)
+                add_special_tokens = not _chat.writes_special_tokens(tokenizer)
+            else:
+                prompt = config["prompt"]
+                add_special_tokens = True
+
+            inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=add_special_tokens).to(device)
+            with torch.no_grad():
+                output_ids = model.generate(**inputs, **generation_options)
+            completion_ids = output_ids[0][inputs["input_ids"].shape[1]:]
+            completion = tokenizer.decode(completion_ids, skip_special_tokens=True)
+
+        return {
+            "completion": str(completion),
+            "model": config["model"],
+            "device": device,
+            "new_tokens": int(completion_ids.shape[0]),
+        }
+    finally:
+        model = None
+        gc.collect()
+        _device.empty_cache(device)

data/python/rpeft_backend/_merge.py

@@ -0,0 +1,37 @@
+import gc
+
+from . import _models
+from ._errors import stage
+
+
+def merge_lora(config):
+    with stage("merge"):
+        from peft import PeftConfig, PeftModel
+        from transformers import AutoModelForCausalLM
+
+        adapter_path = config["adapter_path"]
+        output_dir = config["output_dir"]
+        base_model = config.get("base_model") or PeftConfig.from_pretrained(
+            adapter_path
+        ).base_model_name_or_path
+
+        merged = None
+        try:
+            # Merging is elementwise weight math; CPU is
+            # the memory-safest place for it and speed is irrelevant.
+            base = AutoModelForCausalLM.from_pretrained(base_model, dtype="auto")
+            model = PeftModel.from_pretrained(base, adapter_path)
+            merged = model.merge_and_unload()
+            merged.save_pretrained(output_dir, safe_serialization=True)
+
+            # From the base model, the same tokenizer generation with the adapter uses. The adapter
+            # directory may carry a tokenizer and chat template of its own, which would otherwise
+            # make the merged model behave differently from the adapter it was merged from.
+            tokenizer = _models.load_tokenizer(base_model)
+            tokenizer.save_pretrained(output_dir)
+
+            dtype = str(next(merged.parameters()).dtype).replace("torch.", "")
+            return {"model_path": output_dir, "base_model": base_model, "dtype": dtype}
+        finally:
+            merged = None
+            gc.collect()

data/python/rpeft_backend/_models.py

@@ -0,0 +1,30 @@
+from ._errors import stage
+
+
+def load_tokenizer_and_model(base_model, dtype, device, trust_remote_code=False):
+    with stage("model_load"):
+        import torch
+        from transformers import AutoModelForCausalLM
+
+        tokenizer = load_tokenizer(base_model, trust_remote_code=trust_remote_code)
+        model = AutoModelForCausalLM.from_pretrained(
+            base_model,
+            dtype=getattr(torch, dtype),
+            trust_remote_code=trust_remote_code,
+        )
+        model = model.to(device)
+        model.config.pad_token_id = tokenizer.pad_token_id
+
+        return tokenizer, model
+
+
+def load_tokenizer(base_model, trust_remote_code=False):
+    from transformers import AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        base_model, use_fast=True, trust_remote_code=trust_remote_code
+    )
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    return tokenizer