lora-easy 2.0__tar.gz → 2.2__tar.gz

{lora_easy-2.0 → lora_easy-2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lora-easy
-Version: 2.0
+Version: 2.2
 Summary: A tiny OOP wrapper around PEFT for LoRA fine-tuning of causal LMs.
 Author: Freakwill
 License: MIT

{lora_easy-2.0 → lora_easy-2.2}/lora_easy.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lora-easy
-Version: 2.0
+Version: 2.2
 Summary: A tiny OOP wrapper around PEFT for LoRA fine-tuning of causal LMs.
 Author: Freakwill
 License: MIT

{lora_easy-2.0 → lora_easy-2.2}/lora_ez/agent.py

@@ -29,8 +29,9 @@ class Agent:
     name : str or None
         Display name.  Defaults to ``model.name``.
     description : str or None
-        The agent's persona — forwarded as ``system_prompt`` on every turn.
-        Defaults to ``model.system_prompt``.
+        The agent's persona.  Defaults to ``model.description`` (the immutable
+        identity).  Only forwarded as ``system_prompt`` when it differs from
+        the model's own description.
     """
 
     def __init__(
@@ -51,7 +52,7 @@ class Agent:
     ):
         self.model = model
         self.name = name or model.name
-        self.description = description or model.system_prompt
+        self.description = description or model.description
 
         self.web_enabled = web_enabled
         self.web_allowlist = web_allowlist or []
@@ -81,7 +82,7 @@ class Agent:
         return cls(
             model,
             name=cfg.get("name", model.name),
-            description=cfg.get("description", model.system_prompt),
+            description=cfg.get("description", model.description),
             web_enabled=cfg.get("web_enabled", False),
             web_allowlist=cfg.get("web_allowlist"),
             web_blocklist=cfg.get("web_blocklist"),
@@ -115,9 +116,14 @@ class Agent:
                 f"{prompt}\n\n"
                 f"[Relevant context from tools:\n{context}\n]"
             )
+        # model.description is prepended by the model itself; only forward a
+        # system_prompt that differs from it (avoids duplicating the identity)
+        sp = system_prompt
+        if sp is None and self.description and self.description != self.model.description:
+            sp = self.description
         return self.model.chat(
             prompt, history=history, max_tokens=max_tokens,
-            system_prompt=system_prompt or self.description, **kwargs,
+            system_prompt=sp, **kwargs,
         )
 
     # -- context gathering ---------------------------------------------------
@@ -223,8 +229,11 @@ class Agent:
 
     def chat_session(self, *args, **kwargs):
         from .session import _ChatSession
+        sp = kwargs.pop("system_prompt", None)
+        if sp is None and self.description != self.model.description:
+            sp = self.description
         return _ChatSession(
             self.model, *args,
-            system_prompt=kwargs.pop("system_prompt", self.description),
+            system_prompt=sp,
             **kwargs,
         )

{lora_easy-2.0 → lora_easy-2.2}/lora_ez/model.py

@@ -22,9 +22,8 @@ Use::
         s.run()                    # interactive REPL, /exit to quit
 """
 
-from pathlib import Path
-
 import torch
+from pathlib import Path
 from peft import LoraConfig, get_peft_model, PeftModel
 from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
 
@@ -34,7 +33,7 @@ class LoraModel:
 
     Parameters
     ----------
-    model_id : str
+    id_ : str
         HuggingFace model identifier (e.g. ``"Qwen/Qwen2.5-0.5B-Instruct"``).
     name : str
         Human-readable label, used for default save paths and REPL display.
@@ -43,6 +42,16 @@ class LoraModel:
     system_prompt : str or None
         Prepend a system instruction to every chat turn.  When ``None`` the
         model uses its built-in default (``"You are Qwen …"``).
+    description : str or None
+        Immutable identity/role (e.g. ``"you are a sassy house cat"``).
+        Rendered as its own system message BEFORE ``system_prompt``.
+        Saved with the adapter (``description.txt``) and restored on ``load()``.
+        ``None`` (default) means no description is stored or loaded.
+    save_path : str
+        The path where the adapter is saved.
+    cache_dir : str or None
+        Custom model download cache.  ``None`` uses the default
+        ``~/.cache/huggingface/hub/``.
 
     Key attributes
     --------------
@@ -54,34 +63,46 @@ class LoraModel:
         The active model — ``base`` or ``peft_model`` depending on LoRA state.
     """
 
-    def __init__(self, model_id: str, name: str = "Assistant", device: str = "mps",
-                 system_prompt: str | None = None):
-        self.model_id = model_id
+    def __init__(self, id_: str, name: str = "Assistant", device: str = "mps",
+                 system_prompt: str | None = None, save_path: str | None = None,
+                 cache_dir: str | None = None, description: str | None = None):
+        self.id_ = id_
         self.name = name
         self.system_prompt = system_prompt
-        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+        self._description = description
+        # tokenizer download (cached at cache_dir or ~/.cache/huggingface/hub/)
+        self.tokenizer = AutoTokenizer.from_pretrained(id_, cache_dir=cache_dir)
         self.tokenizer.pad_token = self.tokenizer.eos_token
-        self.base = AutoModelForCausalLM.from_pretrained(model_id, device_map=device, torch_dtype="auto")
+        # model weights download (largest; first run downloads GBs)
+        self.base = AutoModelForCausalLM.from_pretrained(id_, device_map=device, torch_dtype="auto", cache_dir=cache_dir)
         self.peft_model = None
         self.model = self.base
         self._lora_enabled = False
+        self.save_path = save_path or f"{self:l}-lora"
+
+    @property
+    def description(self) -> str | None:
+        """Immutable role/identity string (read-only; set via ``__init__`` or ``load()``)."""
+        return self._description
 
     def __repr__(self) -> str:
         n = sum(p.numel() for p in self.model.parameters())
-        return f"{self.name.title()}: LoraModel('{self.model_id}', {n/1e9:.1f}B params, lora={'yes' if self._has_lora() else 'no'})"
+        return f"{self:t}: LoraModel('{self:i}', {n/1e9:.1f}B params, lora={'yes' if self._has_lora() else 'no'})"
 
     def __str__(self):
         return self.name.title()
 
-    def __format__(self, spec=None):
-        if spec is None:
+    def __format__(self, spec=''):
+        if spec == '':
             return self.name
-        elif spec is 'l':
+        elif spec == 'l':
             return self.name.lower()
-        elif spec is 't':
+        elif spec == 't':
             return self.name
+        elif spec == 'i':
+            return self.id_
         else:
-            raise ValueError('`spec` should be one of None | l | t')
+            raise ValueError('`spec` should be one of [empty] | l | t')
 
     def _has_lora(self) -> bool:
         return self.peft_model is not None
@@ -130,12 +151,19 @@ class LoraModel:
         ``history`` is a list of ``{"role": ..., "content": ...}`` dicts from
         previous turns.  The model sees the full context and can refer back to it.
         ``system_prompt`` overrides ``self.system_prompt`` for this single turn.
+        ``description`` (if set) is always prepended as its own system message.
         """
         self.model.eval()
         sp = system_prompt if system_prompt is not None else self.system_prompt
         messages = list(history or [])
-        if sp and (not messages or messages[0].get("role") != "system"):
-            messages = [{"role": "system", "content": sp}] + messages
+        if not messages or messages[0].get("role") != "system":
+            # description first (immutable identity), then system_prompt (mutable)
+            sys_msgs = []
+            if self.description:
+                sys_msgs.append({"role": "system", "content": self.description})
+            if sp:
+                sys_msgs.append({"role": "system", "content": sp})
+            messages = sys_msgs + messages
         messages.append({"role": "user", "content": prompt})
         fmt = self._format(messages)
         inp = self.tokenizer(fmt, return_tensors="pt").to(self.model.device)
@@ -157,13 +185,12 @@ class LoraModel:
         Usage::
 
             with model.chat_session("./history.json", auto_save=True) as s:
-                s.chat("Hello")
-                s.chat("What do you think?")
+                s> "Hello"
+                s> "What do you think?"
                 s.history   # all turns so far
 
         When the history grows beyond ``max_tokens`` tokens it is automatically
         summarised by the model and replaced with a compact system message.
-        ``system_prompt`` overrides ``self.system_prompt`` for the session.
         """
         from .session import _ChatSession
         return _ChatSession(self, save_path, auto_save, max_tokens,
@@ -171,7 +198,8 @@ class LoraModel:
 
     # -- Training -----------------------------------------------------------
 
-    def train(self, conversations: list[dict], output: bool = False, **kwargs):
+    def train(self, conversations: list[dict], save_checkpoints: bool = False,
+              max_length: int = 256, **kwargs):
         """Fine-tune with LoRA on ShareGPT-format conversations.
 
         Pipeline: render each convo via chat template -> tokenize with pad+trunc ->
@@ -179,30 +207,39 @@ class LoraModel:
         Auto-enables LoRA if not already active.
 
         Args:
-            output: if set, save training logs/checkpoints to ``./lora-output-{name}``.
-                    False (default) produces no output files.
+            save_checkpoints: if set, save a checkpoint per epoch to
+                    ``./lora-output-{name}``.  False (default) produces no files.
+            max_length: tokenizer truncation/padding length.  Raise it when
+                    conversations are long (multi-turn or long replies) —
+                    truncation cuts from the END, which would clip the target
+                    assistant reply.  Default 256.
         """
         if not self.lora_enabled:
             self.enable_lora()
         self.model.config.use_cache = False
 
         texts = [self._format(self._render(c), add_gen=False) for c in conversations]
-        tok = self.tokenizer(texts, truncation=True, padding="max_length", max_length=256)
+        tok = self.tokenizer(texts, truncation=True, padding="max_length",
+                             max_length=max_length)
         dataset = [{"input_ids": inds, "attention_mask": ms,
                "labels": [-100 if m == 0 else i for i, m in zip(inds, ms)]}
               for inds, ms in zip(tok["input_ids"], tok["attention_mask"])]
 
-        output_dir = f"./lora-output-{self.name.lower()}" if output else "./temp_output"
-        defaults = {"epochs": 30, "lr": 3e-4, "per_device_train_batch_size": 4, "logging_steps": 5}
-        merged = defaults | kwargs
+        # friendly defaults; translate to TrainingArguments names below
+        defaults = {"output_dir": f"./lora-output-{self:l}" if save_checkpoints else "./temp_output",
+                    "epochs": 30, "lr": 3e-4,
+                    "per_device_train_batch_size": 4, "logging_steps": 5}
+        kwargs = defaults | kwargs
         args = TrainingArguments(
-            output_dir=output_dir, **merged,
-            save_strategy="no" if output is None else "epoch",
+            num_train_epochs=kwargs.pop("epochs"),
+            learning_rate=kwargs.pop("lr"),
+            **kwargs,
+            save_strategy="no" if not save_checkpoints else "epoch",
             report_to="none",
         )
         Trainer(model=self.model, args=args, train_dataset=dataset, processing_class=self.tokenizer).train()
 
-        if not output:
+        if not save_checkpoints:
             import shutil
             shutil.rmtree("./temp_output", ignore_errors=True)
 
@@ -218,24 +255,43 @@ class LoraModel:
     def from_yaml(cls, path: str):
         """Create a LoraModel from a YAML config file.
 
-        Expected keys: ``model_id``, ``name``, ``system_prompt``,
+        Expected keys: ``id_``, ``name``, ``system_prompt``,
         ``peft_path`` (optional — auto-loads the adapter).
         """
         import yaml
+        from pathlib import Path
         cfg = yaml.safe_load(Path(path).read_text())
-        m = cls(cfg["model_id"], name=cfg.get("name", "Assistant"),
-                 system_prompt=cfg.get("system_prompt"))
+        if "id_" in cfg:
+            id_ = cfg["id_"]
+        elif "id" in cfg:
+            id_ = cfg["id"]
+        else:
+            raise KeyError("Not provide the key `id_` | `id`.")
+        m = cls(id_=id_, name=cfg.get("name", "Assistant"),
+                 system_prompt=cfg.get("system_prompt"),
+                 description=cfg.get("description"))
         if peft := cfg.get("peft_path"):
             m.load(peft)
         return m
 
     def save(self, path: str | None = None):
-        path = f"lora-{self:l}" if path is None else path
+        # save the adapter
+        path = path or self.save_path
         self.model.save_pretrained(path)
         self.tokenizer.save_pretrained(path)
-        print(f"The adapter is saved in `{path}`.")
+        if self.description is not None:
+            (Path(path) / "description.txt").write_text(self.description)
 
     def load(self, path: str | None = None):
-        path = f"lora-{self:l}" if path is None else path
-        self.peft_model = PeftModel.from_pretrained(self.base, path)
-        # self.enable_lora()
+        # load the adapter from a LOCAL directory
+        p = Path(path or self.save_path)
+        if not p.is_dir():
+            raise FileNotFoundError(
+                f"""Adapter directory not found: '{p}'. 
+If it lives elsewhere pass the path explicitly, e.g. m.load('/path/to/ivanka-lora').""")
+        self.peft_model = PeftModel.from_pretrained(self.base, str(p),
+                                                    is_trainable=True)
+        # restore description if it was saved alongside the adapter
+        desc_file = p / "description.txt"
+        if desc_file.exists():
+            self._description = desc_file.read_text()

{lora_easy-2.0 → lora_easy-2.2}/lora_ez/session.py

@@ -73,11 +73,15 @@ class _ChatSession:
         """
         self._maybe_summarize(prompt)
         self.history.append({"role": "user", "content": prompt})
-        reply = self.model.chat(prompt, history=self.history[:-1],
+        reply = self.model.chat(prompt, history=self.history[:-1], 
                                 system_prompt=self.system_prompt, **kwargs)
         self.history.append({"role": "assistant", "content": reply})
         return reply
 
+    def __gt__(self, prompt: str) -> str:
+        """``s > "prompt"`` is shorthand for ``s.chat("prompt")``."""
+        return self.chat(prompt)
+
     def run(self):
         """Interactive REPL — type messages at a prompt, ``/exit`` to quit."""
         print(f"Chat session started.  Type /exit to quit.\n"
@@ -127,10 +131,14 @@ class _ChatSession:
         raw = self.model.tokenizer.decode(out[0], skip_special_tokens=True)
         summary = raw.rpartition("assistant\n")[-1].strip()
 
-        # preserve system prompt if set, merge summary into it
+        # preserve description (immutable) as its own system message,
+        # then merge summary into the system_prompt (mutable)
+        self.history = []
+        if self.model.description:
+            self.history.append({"role": "system", "content": self.model.description})
         if self.system_prompt:
-            self.history = [{"role": "system",
-                                 "content": f"{self.system_prompt}\n\n[Summary of previous turns: {summary}]"}]
+            self.history.append({"role": "system",
+                                 "content": f"{self.system_prompt}\n\n[Summary of previous turns: {summary}]"})
         else:
-            self.history = [{"role": "system",
-                                 "content": f"Previous conversation summary: {summary}"}]
+            self.history.append({"role": "system",
+                                 "content": f"Previous conversation summary: {summary}"})

{lora_easy-2.0 → lora_easy-2.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "lora-easy"
-version = "2.0"
+version = "2.2"
 description = "A tiny OOP wrapper around PEFT for LoRA fine-tuning of causal LMs."
 readme = "README.md"
 requires-python = ">=3.9"