erictransformer 0.0.1__py3-none-any.whl

erictransformer/__init__.py

@@ -0,0 +1,44 @@
+from erictransformer.eric_tasks import (
+    CHATCallArgs,
+    CHATResult,
+    CHATStreamResult,
+    CHATTokArgs,
+    GENCallArgs,
+    GENResult,
+    GENTokArgs,
+    EricChat,
+    EricGeneration,
+    EricTextClassification,
+    EricTextToText,
+    TCCallArgs,
+    TCResult,
+    TCTokArgs,
+    TTCallArgs,
+    TTResult,
+    TTTokArgs
+)
+
+from erictransformer.args import EricTrainArgs, EricEvalArgs
+from erictransformer.eric_transformer import EricTransformer
+from erictransformer.loops import EvalResult, TrainResult
+
+__all__ = []
+
+try:
+    from mlx_lm import load, stream_generate
+    from mlx_lm.sample_utils import make_sampler
+    from mlx_lm.utils import save_model
+    mlx_enabled = True
+except ImportError as err:
+    mlx_enabled = False
+
+if mlx_enabled:
+    try:
+        from .eric_tasks.eric_chat_mlx import EricChatMLX
+    except Exception:
+        # Optional: log or ignore
+        pass
+    else:
+        __all__.append("EricChatMLX")
+
+name = "erictransformer"

erictransformer/args/__init__.py

@@ -0,0 +1,7 @@
+from erictransformer.args.eric_args import (
+    CallResult,
+    CallArgs,
+    EricEvalArgs,
+    TokArgs,
+    EricTrainArgs,
+)

erictransformer/args/eric_args.py

@@ -0,0 +1,50 @@
+from dataclasses import dataclass
+
+@dataclass(kw_only=True)
+class EricTrainArgs:
+    # Learning parameters
+    lr: float = 2e-5
+    bs: int = 1
+    eval_bs: int = 0  # when 0 uses bs
+    epochs: int = 1
+    gas: int = 1
+    optim: str = "adamw"  # options adamw and sgd
+    lr_sched: str = "constant"  # other option: warmup_then_decay
+
+    # Action steps
+    eval_steps: int = 256  # if 0 no evaluating will be done
+    log_steps: int = 256  # if 0 no logging will be done
+    checkpoint_steps: int = -1  # if -1 no checkpointing will be done
+    save_best: bool = False  # saves the model with the lowest eval loss
+
+    # Misc
+    out_dir: str = "eric_transformer/"
+    run_name: str = ""
+    seed: int = 42
+
+
+@dataclass(kw_only=True)
+class TokArgs:
+    bs: int = 1024
+    max_cases: int = -1
+    shards: int = 1
+    procs: int = -1
+
+
+
+@dataclass(kw_only=True)
+class EricEvalArgs:
+    bs: int = 1
+    out_dir: str = "eric_transformer/"
+    run_name: str = ""
+    seed: int = 42
+
+
+@dataclass(kw_only=True)
+class CallArgs:
+    pass
+
+
+@dataclass(kw_only=True)
+class CallResult:
+    pass

erictransformer/eric_tasks/__init__.py

@@ -0,0 +1,47 @@
+from erictransformer.eric_tasks.args import (
+    CHATCallArgs,
+    CHATTokArgs,
+    GENCallArgs,
+    GENTokArgs,
+    TCCallArgs,
+    TCTokArgs,
+    TTCallArgs,
+    TTTokArgs,
+)
+from erictransformer.eric_tasks.chat_stream_handlers import CHATStreamResult
+from erictransformer.eric_tasks.eric_chat import EricChat
+from erictransformer.eric_tasks.eric_generation import EricGeneration
+from erictransformer.eric_tasks.eric_text_classification import (
+    EricTextClassification,
+)
+from erictransformer.eric_tasks.eric_text_to_text import (
+    EricTextToText,
+    TTStreamResult,
+)
+from erictransformer.eric_tasks.results import (
+    CHATResult,
+    GENResult,
+    TCResult,
+    TTResult
+)
+
+__all__ = []
+
+
+try:
+    from mlx_lm import load, stream_generate
+    from mlx_lm.sample_utils import make_sampler
+    from mlx_lm.utils import save_model
+    mlx_enabled = True
+except ImportError as err:
+    mlx_enabled = False
+
+
+if mlx_enabled:
+    try:
+        from .eric_chat_mlx import EricChatMLX
+    except Exception:
+        # Optional: log or ignore
+        pass
+    else:
+        __all__.append("EricChatMLX")

erictransformer/eric_tasks/args/__init__.py

@@ -0,0 +1,16 @@
+from erictransformer.eric_tasks.args.eric_chat_args import (
+    CHATCallArgs,
+    CHATTokArgs,
+)
+from erictransformer.eric_tasks.args.eric_generation_args import (
+    GENCallArgs,
+    GENTokArgs,
+)
+from erictransformer.eric_tasks.args.eric_text_classification_args import (
+    TCCallArgs,
+    TCTokArgs,
+)
+from erictransformer.eric_tasks.args.eric_text_to_text_args import (
+    TTCallArgs,
+    TTTokArgs,
+)

erictransformer/eric_tasks/args/eric_chat_args.py

@@ -0,0 +1,21 @@
+from dataclasses import dataclass, field
+from ericsearch import SearchCallArgs
+
+from erictransformer.args import CallArgs, TokArgs
+
+
+@dataclass(kw_only=True)
+class CHATCallArgs(CallArgs):
+    min_len: int = 1
+    max_len: int = 4096
+    temp: float = 0.8
+    top_k: int = 32
+    top_p: float = 0.6
+
+    # dataset args
+    search_args: SearchCallArgs = field(default_factory=SearchCallArgs)
+
+
+@dataclass(kw_only=True)
+class CHATTokArgs(TokArgs):
+    max_len: int = -1

erictransformer/eric_tasks/args/eric_generation_args.py

@@ -0,0 +1,20 @@
+from dataclasses import dataclass
+
+from erictransformer.args import CallArgs, TokArgs
+
+
+@dataclass(kw_only=True)
+class GENTokArgs(TokArgs):
+    max_len: int = 1024
+
+
+@dataclass(kw_only=True)
+class GENCallArgs(CallArgs):  # ← new canonical name
+    min_len: int = 1
+    max_len: int = 1024
+    temp: float = 0.8
+    top_k: int = 32
+    top_p: float = 0.6
+
+
+

erictransformer/eric_tasks/args/eric_text_classification_args.py

@@ -0,0 +1,13 @@
+from dataclasses import dataclass
+
+from erictransformer.args import CallArgs, TokArgs
+
+
+@dataclass(kw_only=True)
+class TCTokArgs(TokArgs):
+    max_len: int = -1
+
+
+@dataclass(kw_only=True)
+class TCCallArgs(CallArgs):
+    pass

erictransformer/eric_tasks/args/eric_text_to_text_args.py

@@ -0,0 +1,18 @@
+from dataclasses import dataclass
+
+from erictransformer.args import CallArgs, TokArgs
+
+
+@dataclass(kw_only=True)
+class TTCallArgs(CallArgs):
+    min_len: int = 1
+    max_len: int = 1024
+    temp: float = 0.8
+    top_k: int = 32
+    top_p: float = 0.6
+
+
+@dataclass(kw_only=True)
+class TTTokArgs(TokArgs):
+    max_in_len: int = -1
+    max_out_len: int = -1

erictransformer/eric_tasks/chat_stream_handlers/__init__.py

@@ -0,0 +1,6 @@
+from erictransformer.eric_tasks.chat_stream_handlers.args import CHATStreamResult
+from erictransformer.eric_tasks.chat_stream_handlers.default import (
+    DefaultStreamHandler,
+)
+from erictransformer.eric_tasks.chat_stream_handlers.gpt_oss import GPTOSSSMHandler
+from erictransformer.eric_tasks.chat_stream_handlers.smol import SmolStreamHandler

erictransformer/eric_tasks/chat_stream_handlers/args.py

@@ -0,0 +1,13 @@
+from dataclasses import dataclass
+from typing import Optional
+
+@dataclass(kw_only=True)
+class CHATStreamResult:
+    text: str  # what user sees
+    marker: str  # the marker e.g text, special, tool
+    payload: Optional[dict]
+
+
+@dataclass(kw_only=True)
+class MarkerStrings:
+    pass

erictransformer/eric_tasks/chat_stream_handlers/default.py

@@ -0,0 +1,19 @@
+from typing import Union
+
+from transformers import PreTrainedTokenizer
+
+from erictransformer.eric_tasks.chat_stream_handlers.args import CHATStreamResult
+from erictransformer.eric_tasks.chat_stream_handlers.stream_handler import (
+    StreamHandler,
+)
+
+
+class DefaultStreamHandler(StreamHandler):
+    def __init__(self, tokenizer: PreTrainedTokenizer):
+        super().__init__(tokenizer=tokenizer)
+
+    def step(self, token_str: str) -> Union[None, CHATStreamResult]:
+        if not token_str:
+            return None
+
+        return CHATStreamResult(text=token_str, marker="text", payload={})

erictransformer/eric_tasks/chat_stream_handlers/gpt_oss.py

@@ -0,0 +1,147 @@
+from dataclasses import dataclass
+from typing import Union
+import logging
+
+from transformers import PreTrainedTokenizer
+
+from erictransformer.eric_tasks.chat_stream_handlers.args import (
+    CHATStreamResult,
+    MarkerStrings,
+)
+from erictransformer.eric_tasks.chat_stream_handlers.stream_handler import (
+    StreamHandler,
+)
+
+
+@dataclass(kw_only=True)
+class GPTOSSMarkerStrings(MarkerStrings):
+    start: str = "<|start|>"
+    end: str = "<|end|>"
+    message: str = "<|message|>"
+    channel: str = "<|channel|>"
+    constrain: str = "<|constrain|>"
+    return_token: str = "<|return|>"
+    call: str = "<|call|>"
+
+
+class GPTOSSSMHandler(StreamHandler):
+    def __init__(self, tokenizer: PreTrainedTokenizer, logger: logging.Logger):
+        self.markers = GPTOSSMarkerStrings()
+        self.current_channel = ""
+        self.in_thinking = False
+        self.in_tool = False
+        self.change_channel = False
+        self.just_received_message = False
+        self.start_just_happened = False
+        self.tool_strings = []
+        self.logger = logger
+        super().__init__(tokenizer=tokenizer)
+
+    def _reset_state(self):
+        self.current_channel = ""
+        self.in_thinking = False
+        self.in_tool = False
+        self.change_channel = False
+        self.just_received_message = False
+        self.start_just_happened = False
+        self.tool_strings = []
+
+    def step(self, token_str: str) -> Union[None, CHATStreamResult]:
+        if not token_str:
+            return None
+
+        stripped_token_str = token_str.strip()
+
+        self.in_thinking = False
+        #### SPECIAL TOKENS ####
+
+        # Case: <|start|>
+        if stripped_token_str == self.markers.start:
+            self._reset_state()
+            self.start_just_happened = True
+            return CHATStreamResult(
+                text=stripped_token_str, marker="special", payload={}
+            )
+
+        # Case: <|channel|>
+        if stripped_token_str == self.markers.channel:
+            self.change_channel = True
+            # next round
+            return CHATStreamResult(
+                text=stripped_token_str, marker="special", payload={}
+            )
+
+        # Case: <|message|>
+        if stripped_token_str == self.markers.message:
+
+            return CHATStreamResult(
+                text=stripped_token_str, marker="special", payload={}
+            )
+
+        # Case: <|return|>
+        if stripped_token_str == self.markers.return_token:
+            # do nothing. Streaming is over.
+            self._reset_state()
+            return CHATStreamResult(
+                text=stripped_token_str, marker="special", payload={}
+            )
+
+        # Case: <|constrain|>
+        if stripped_token_str == self.markers.constrain:
+            # do nothing. we always assume json
+            return CHATStreamResult(
+                text=stripped_token_str, marker="special", payload={}
+            )
+
+        # Case: <|end|>
+        if stripped_token_str == self.markers.end:
+            temp_current_channel = self.current_channel
+            self._reset_state()
+            if temp_current_channel == "analysis":
+                return CHATStreamResult(
+                    text=stripped_token_str, marker="think_end", payload={}
+                )
+            return CHATStreamResult(
+                text=stripped_token_str, marker="special", payload={}
+            )
+
+        #### NON SPECIAL TOKENS ####
+
+        if self.start_just_happened:
+            # we don't do anything with this header token
+            self.start_just_happened = False
+
+            if stripped_token_str in {"assistant", "user", "system"}:
+                return CHATStreamResult(text=token_str, marker="special", payload={})
+
+        if self.change_channel:
+            self.change_channel = False
+
+            if self.current_channel == "commentary":
+                self.logger.warning("The 'commentary' channel is not supported. Falling back to 'analysis'. ")
+                self.current_channel = "analysis"
+
+            if stripped_token_str not in ("analysis", "final"):
+                # fall back to final
+                self.logger.warning(f"Unexpected channel '{stripped_token_str}'. Falling back to 'final'.")
+                self.current_channel = "final"
+            else:
+                self.current_channel = stripped_token_str
+
+            if stripped_token_str == "analysis":
+                return CHATStreamResult(
+                    text=stripped_token_str, marker="think_start", payload={}
+                )
+
+            return CHATStreamResult(
+                text=stripped_token_str, marker="special", payload={}
+            )
+        if self.current_channel == "analysis":
+            # just return the text
+            return CHATStreamResult(text=token_str, marker="thinking", payload={})
+
+        if self.current_channel == "final":
+            # just return the text
+            return CHATStreamResult(text=token_str, marker="text", payload={})
+
+        self._reset_state()

erictransformer/eric_tasks/chat_stream_handlers/smol.py

@@ -0,0 +1,81 @@
+from dataclasses import dataclass
+from typing import Union
+import logging
+
+from transformers import PreTrainedTokenizer
+
+from erictransformer.eric_tasks.chat_stream_handlers.args import (
+    CHATStreamResult,
+    MarkerStrings,
+)
+from erictransformer.eric_tasks.chat_stream_handlers.stream_handler import (
+    StreamHandler,
+)
+
+
+@dataclass(kw_only=True)
+class SMOLMarkerStrings(MarkerStrings):
+    think_start: str
+    think_end: str
+    tool_start: str
+    tool_end: str
+
+
+class SmolStreamHandler(StreamHandler):
+    def __init__(self, tokenizer: PreTrainedTokenizer,  logger: logging.Logger):
+        self.markers: SMOLMarkerStrings = SMOLMarkerStrings(
+            think_start="<think>",
+            think_end="</think>",
+            tool_start="<tool_call>", # not supported
+            tool_end="</tool_call>", # not supported
+        )
+
+        self.marker_key_tuple = (
+            self.markers.think_start,
+            self.markers.think_end,
+            self.markers.tool_start,
+            self.markers.tool_end,
+        )
+
+        self.in_thinking = False
+
+        self.logger = logger
+
+        super().__init__(tokenizer=tokenizer)
+
+    def step(self, token_str: str) -> Union[None, CHATStreamResult]:
+        if not token_str:
+            return None
+
+        stripped_token_str = token_str.strip()
+
+        if stripped_token_str not in self.marker_key_tuple:
+            if self.in_thinking:
+                return CHATStreamResult(text=token_str, marker="thinking", payload={})
+
+            if stripped_token_str in self.special_tokens:
+                return CHATStreamResult(
+                    text=stripped_token_str, marker="special", payload={})
+
+            else:
+                if token_str.endswith(self.eos_token):
+                    token_str = token_str[: -len(self.eos_token)]
+                return CHATStreamResult(text=token_str, marker="text", payload={})
+
+        elif stripped_token_str == self.markers.think_start:
+            self.in_thinking = True
+            return CHATStreamResult(
+                text=stripped_token_str, marker="think_start", payload={}
+            )
+
+        elif stripped_token_str == self.markers.think_end:
+            self.in_thinking = False
+            return CHATStreamResult(
+                text=stripped_token_str, marker="think_end", payload={}
+            )
+
+        elif stripped_token_str in (self.markers.tool_start, self.markers.tool_end):
+            self.logger.warning("Tool calling is not supported but a tool token was generated")
+            return None
+
+        return None

erictransformer/eric_tasks/chat_stream_handlers/stream_handler.py

@@ -0,0 +1,17 @@
+from abc import ABC, abstractmethod
+from typing import Union
+
+from transformers import PreTrainedTokenizer
+
+from erictransformer.eric_tasks.chat_stream_handlers.args import CHATStreamResult
+
+
+class StreamHandler(ABC):
+    def __init__(self, tokenizer: PreTrainedTokenizer):
+        self.tokenizer = tokenizer
+        self.special_tokens = tokenizer.all_special_tokens
+        self.eos_token = self.tokenizer.eos_token
+
+    @abstractmethod
+    def step(self, token_str: str) -> Union[None, CHATStreamResult]:
+        pass

erictransformer/eric_tasks/chat_templates/__init__.py

@@ -0,0 +1 @@
+from erictransformer.eric_tasks.chat_templates.convert import map_chat_roles

erictransformer/eric_tasks/chat_templates/convert.py

@@ -0,0 +1,67 @@
+from erictransformer.exceptions import EricInputError
+
+
+def map_chat_roles(messages: list, model_name: str, model_type: str):
+    role_map = get_role_map_for_model(model_name)
+    mapped_messages = []
+
+    if model_type == "smollm3":
+        mapped_messages.append({"role": "system", "content": "/think"})
+
+    for m in messages:
+        original_role = m.get("role")
+        mapped_role = role_map.get(original_role)
+
+        if mapped_role is None:
+            raise EricInputError(
+                f"Unsupported role '{original_role}' for model '{model_name}'"
+            )
+
+        mapped_messages.append({"role": mapped_role, "content": m["content"]})
+
+    return mapped_messages
+
+
+def get_role_map_for_model(model_name: str):
+    name = model_name.lower()
+    # Standard/default families
+    if any(
+        k in name
+        for k in (
+            "granite",
+            "smol",
+            "llama",
+            "mistral",
+            "mixtral",
+            "hermes",
+            "dialogpt",
+            "openchat",
+            "chatml",
+        )
+    ):
+        return {
+            "system": "system",
+            "user": "user",
+            "assistant": "assistant",
+            "tool": "tool",
+        }
+
+    if "gemma" in name:
+        return {"system": "user", "user": "user", "assistant": "model", "tool": "tool"}
+
+    # Families where 'system' should be treated as 'user'
+    if "falcon" in name or "vicuna" in name or "alpaca" in name:
+        return {
+            "system": "user",
+            "user": "user",
+            "assistant": "assistant",
+            "tool": "tool",
+        }
+
+    # Fallback
+    return {
+        "system": "system",
+        "user": "user",
+        "assistant": "assistant",
+        "tool": "tool",
+    }