Instructions to use IvmeLabs/Ivme-Conversate-P-v1-Base with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use IvmeLabs/Ivme-Conversate-P-v1-Base with Transformers:

# Use a pipeline as a high-level helper
from transformers import pipeline

pipe = pipeline("text-generation", model="IvmeLabs/Ivme-Conversate-P-v1-Base", trust_remote_code=True)

# pip install -U transformers accelerate
# Load model directly
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("IvmeLabs/Ivme-Conversate-P-v1-Base", trust_remote_code=True, device_map="auto")

Notebooks
Google Colab
Kaggle
Local Apps Settings

vLLM

How to use IvmeLabs/Ivme-Conversate-P-v1-Base with vLLM:

Install from pip and serve model

# Install vLLM from pip:
pip install vllm
# Start the vLLM server:
vllm serve "IvmeLabs/Ivme-Conversate-P-v1-Base"
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:8000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "IvmeLabs/Ivme-Conversate-P-v1-Base",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Use Docker

docker model run hf.co/IvmeLabs/Ivme-Conversate-P-v1-Base

SGLang

How to use IvmeLabs/Ivme-Conversate-P-v1-Base with SGLang:

Install from pip and serve model

# Install SGLang from pip:
pip install sglang
# Start the SGLang server:
python3 -m sglang.launch_server \
    --model-path "IvmeLabs/Ivme-Conversate-P-v1-Base" \
    --host 0.0.0.0 \
    --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "IvmeLabs/Ivme-Conversate-P-v1-Base",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Use Docker images

docker run --gpus all \
    --shm-size 32g \
    -p 30000:30000 \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    --env "HF_TOKEN=<secret>" \
    --ipc=host \
    lmsysorg/sglang:latest \
    python3 -m sglang.launch_server \
        --model-path "IvmeLabs/Ivme-Conversate-P-v1-Base" \
        --host 0.0.0.0 \
        --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "IvmeLabs/Ivme-Conversate-P-v1-Base",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Docker Model Runner
How to use IvmeLabs/Ivme-Conversate-P-v1-Base with Docker Model Runner:
```
docker model run hf.co/IvmeLabs/Ivme-Conversate-P-v1-Base
```

Ivme-Conversate-P-v1-Base / modeling_ivme_conversate.py

ereniko

Release İvme-Conversate-P-v1

29f0d05 verified 22 days ago

Raw History Blame Contribute Delete

12.6 kB

"""
modeling_ivme_conversate.py
HuggingFace transformers-compatible model wrapper for Ivme-Conversate-v3.
Wraps the exact model architecture used during training/inference
(RoPE, QK-norm, XSA, SwiGLU, tied embeddings) so it can be loaded via
AutoModelForCausalLM.from_pretrained(..., trust_remote_code=True) and
used by any tool expecting a standard transformers CausalLM interface -
including external benchmark runners like ArithMark's official script,
which load models this way rather than via custom code.
The underlying architecture code (RMSNorm, Attention/XSA, RoPE, SwiGLU)
is copied verbatim from the verified, tested implementation used for
training and the standalone inference script - including its
dtype-correctness fixes - so results are identical to what training
and infer_ivme_v3.py already produce. This file only adds the
transformers-facing wrapper (config plumbing, forward() output format,
generate() compatibility) around that unchanged core.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast

try:
    from .configuration_ivme_conversate import IvmeConversateConfig
except ImportError:
    # Fallback for direct/standalone imports (e.g. a verification script
    # importing this file directly rather than through transformers'
    # trust_remote_code package-loading machinery, which does support
    # the relative import above).
    from configuration_ivme_conversate import IvmeConversateConfig


# ============================================================================
# Core architecture - copied verbatim from the verified implementation.
# ============================================================================

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        dtype = x.dtype
        x = x.float()
        norm = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return norm.to(dtype) * self.weight.to(dtype)


def precompute_rope(head_dim, max_seq_len, theta=10000.0, device=None):
    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
    t = torch.arange(max_seq_len, device=device).float()
    freqs = torch.outer(t, inv_freq)
    cos = torch.cos(freqs)
    sin = torch.sin(freqs)
    return cos, sin


def apply_rope(x, cos, sin):
    d = x.shape[-1]
    x1, x2 = x[..., : d // 2], x[..., d // 2:]
    cos = cos[None, None, :, :].to(x.dtype)
    sin = sin[None, None, :, :].to(x.dtype)
    rotated = torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
    return rotated


class Attention(nn.Module):
    """Multi-head self-attention with QK-Norm + RoPE + XSA."""

    def __init__(self, cfg):
        super().__init__()
        self.n_heads = cfg.n_heads
        self.head_dim = cfg.d_model // cfg.n_heads
        assert self.head_dim * cfg.n_heads == cfg.d_model

        self.wq = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
        self.wk = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
        self.wv = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
        self.wo = nn.Linear(cfg.d_model, cfg.d_model, bias=False)

        self.q_norm = RMSNorm(self.head_dim, eps=cfg.qknorm_eps)
        self.k_norm = RMSNorm(self.head_dim, eps=cfg.qknorm_eps)

        self.xsa_eps = 1e-6

    def forward(self, x, cos, sin):
        B, T, D = x.shape
        q = self.wq(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.wk(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        v = self.wv(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)

        q = self.q_norm(q)
        k = self.k_norm(k)

        q = apply_rope(q, cos[:T], sin[:T])
        k = apply_rope(k, cos[:T], sin[:T])

        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)

        # --- XSA: z_i = y_i - (y_i . v_i)/(||v_i||^2+eps) * v_i, per token/head ---
        dot = (y * v).sum(dim=-1, keepdim=True)
        vnorm_sq = (v * v).sum(dim=-1, keepdim=True)
        proj_coef = dot / (vnorm_sq + self.xsa_eps)
        z = y - proj_coef * v

        z = z.transpose(1, 2).contiguous().view(B, T, D)
        return self.wo(z)


class SwiGLU(nn.Module):
    def __init__(self, d_model, d_ffn):
        super().__init__()
        self.w_gate = nn.Linear(d_model, d_ffn, bias=False)
        self.w_up = nn.Linear(d_model, d_ffn, bias=False)
        self.w_down = nn.Linear(d_ffn, d_model, bias=False)

    def forward(self, x):
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))


class Block(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.attn_norm = RMSNorm(cfg.d_model, eps=cfg.rmsnorm_eps)
        self.attn = Attention(cfg)
        self.ffn_norm = RMSNorm(cfg.d_model, eps=cfg.rmsnorm_eps)
        self.ffn = SwiGLU(cfg.d_model, cfg.d_ffn)

    def forward(self, x, cos, sin):
        x = x + self.attn(self.attn_norm(x), cos, sin)
        x = x + self.ffn(self.ffn_norm(x))
        return x


# ============================================================================
# transformers-facing wrapper
# ============================================================================

class IvmeConversateModel(PreTrainedModel):
    """Base model (no LM head) - exposed mainly so AutoModel works too;
    most users will want IvmeConversateForCausalLM below."""
    config_class = IvmeConversateConfig

    def __init__(self, config):
        super().__init__(config)
        self.tok_embed = nn.Embedding(config.vocab_size, config.d_model)
        self.blocks = nn.ModuleList([Block(config) for _ in range(config.n_layers)])
        self.final_norm = RMSNorm(config.d_model, eps=config.rmsnorm_eps)
        # NOTE: RoPE cos/sin are deliberately NOT cached as a
        # register_buffer here. transformers v5.x has a known, currently
        # open bug (confirmed via reproduction and matching upstream
        # reports: huggingface/transformers#43950, #44534) where
        # non-persistent buffers (persistent=False) get corrupted with
        # uninitialized/meta-device garbage memory during
        # from_pretrained()'s fast weight-materialization path - this
        # silently produced NaN/garbage RoPE values and completely wrong
        # model output when loaded via AutoModelForCausalLM, even though
        # the actual model weights loaded correctly. Recomputing cos/sin
        # fresh in forward() is cheap (a handful of trig ops) and
        # sidesteps this bug class entirely rather than depending on a
        # library fix.
        self._head_dim = config.d_model // config.n_heads
        self._rope_theta = config.rope_theta
        self._context_len = config.context_len

        self.post_init()

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, input_ids, **kwargs):
        # Defensive truncation: keep only the most recent context_len
        # tokens. This is a second safety layer beyond exposing
        # max_position_embeddings in the config (see
        # IvmeConversateConfig) - if any caller still manages to pass a
        # longer sequence, truncate cleanly to the trailing window
        # rather than crash inside RoPE with a cryptic tensor-shape
        # mismatch (confirmed reproducible: a 2048-token input against
        # this model's 1024 context previously crashed deep inside
        # apply_rope with no indication of the real cause).
        if input_ids.shape[1] > self._context_len:
            input_ids = input_ids[:, -self._context_len:]
        x = self.tok_embed(input_ids)
        cos, sin = precompute_rope(self._head_dim, self._context_len,
                                    theta=self._rope_theta, device=x.device)
        for block in self.blocks:
            x = block(x, cos, sin)
        x = self.final_norm(x)
        return x


class IvmeConversateForCausalLM(PreTrainedModel, GenerationMixin):
    """Standard transformers CausalLM interface: forward() returns
    CausalLMOutputWithPast(logits=...), supports labels for loss
    computation, and is compatible with model.generate() and any tool
    (like ArithMark's benchmark runner) expecting a normal
    AutoModelForCausalLM-loadable model."""

    config_class = IvmeConversateConfig
    # Tells transformers' save_pretrained/tie_weights machinery about the
    # tied embedding explicitly (required in transformers v5+; without
    # this, save_pretrained raises "shared tensors... not properly
    # defined" even though the weights ARE correctly tied via the
    # assignment in __init__ below - confirmed by reproducing the error
    # and fixing it with this attribute).
    _tied_weights_keys = {"lm_head.weight": "model.tok_embed.weight"}

    def __init__(self, config):
        super().__init__(config)
        self.model = IvmeConversateModel(config)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)

        if config.tie_embeddings:
            self.lm_head.weight = self.model.tok_embed.weight

        self.post_init()

    def get_input_embeddings(self):
        return self.model.tok_embed

    def set_input_embeddings(self, value):
        self.model.tok_embed = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
        # attention_mask accepted for interface compatibility but not
        # used - this model is trained/used causal-only with no padding
        # within a sequence (see ShardedTokenDataset / infer scripts).
        #
        # If input_ids exceeds context_len, self.model() truncates it
        # internally (see IvmeConversateModel.forward). labels must be
        # truncated identically here, BEFORE the loss is computed -
        # otherwise a caller passing both an over-length input_ids and
        # matching labels (e.g. a perplexity/loglikelihood task) would
        # hit a shape mismatch between logits and labels, or worse, a
        # silent misalignment if the shapes happened to coincide.
        context_len = self.config.max_position_embeddings
        if input_ids.shape[1] > context_len and labels is not None:
            labels = labels[:, -context_len:]

        hidden = self.model(input_ids)
        logits = self.lm_head(hidden)

        loss = None
        if labels is not None:
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)), labels.view(-1), ignore_index=-100
            )

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
        )

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        # No KV-cache implementation - recomputes the full sequence each
        # step. Correct but not optimized for long generations; adequate
        # for benchmark scoring (log-likelihood over short continuations)
        # and interactive testing.
        return {"input_ids": input_ids}

    def generate(self, *args, **kwargs):
        # Force use_cache=False unconditionally. This model has no cache
        # implementation, and transformers' generate() otherwise tries to
        # construct a DynamicCache using standard config field names
        # (num_hidden_layers, etc.) that this config doesn't define,
        # crashing regardless of the config-level use_cache default set
        # in IvmeConversateConfig (confirmed empirically: that default
        # does not reliably reach generate()'s cache-setup path in all
        # transformers versions). Overriding here guarantees correct
        # behavior for any caller - including external tools like
        # ArithMark's benchmark runner - with zero special knowledge
        # required on their end.
        kwargs["use_cache"] = False
        return super().generate(*args, **kwargs)

    def num_params(self, non_embedding=False):
        n = sum(p.numel() for p in self.parameters())
        if non_embedding:
            n -= self.model.tok_embed.weight.numel()
        return n

	"""
	modeling_ivme_conversate.py

	HuggingFace transformers-compatible model wrapper for Ivme-Conversate-v3.
	Wraps the exact model architecture used during training/inference
	(RoPE, QK-norm, XSA, SwiGLU, tied embeddings) so it can be loaded via
	AutoModelForCausalLM.from_pretrained(..., trust_remote_code=True) and
	used by any tool expecting a standard transformers CausalLM interface -
	including external benchmark runners like ArithMark's official script,
	which load models this way rather than via custom code.

	The underlying architecture code (RMSNorm, Attention/XSA, RoPE, SwiGLU)
	is copied verbatim from the verified, tested implementation used for
	training and the standalone inference script - including its
	dtype-correctness fixes - so results are identical to what training
	and infer_ivme_v3.py already produce. This file only adds the
	transformers-facing wrapper (config plumbing, forward() output format,
	generate() compatibility) around that unchanged core.
	"""

	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from transformers import PreTrainedModel
	from transformers.generation import GenerationMixin
	from transformers.modeling_outputs import CausalLMOutputWithPast

	try:
	from .configuration_ivme_conversate import IvmeConversateConfig
	except ImportError:
	# Fallback for direct/standalone imports (e.g. a verification script
	# importing this file directly rather than through transformers'
	# trust_remote_code package-loading machinery, which does support
	# the relative import above).
	from configuration_ivme_conversate import IvmeConversateConfig


	# ============================================================================
	# Core architecture - copied verbatim from the verified implementation.
	# ============================================================================

	class RMSNorm(nn.Module):
	def __init__(self, dim, eps=1e-5):
	super().__init__()
	self.eps = eps
	self.weight = nn.Parameter(torch.ones(dim))

	def forward(self, x):
	dtype = x.dtype
	x = x.float()
	norm = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
	return norm.to(dtype) * self.weight.to(dtype)


	def precompute_rope(head_dim, max_seq_len, theta=10000.0, device=None):
	inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
	t = torch.arange(max_seq_len, device=device).float()
	freqs = torch.outer(t, inv_freq)
	cos = torch.cos(freqs)
	sin = torch.sin(freqs)
	return cos, sin


	def apply_rope(x, cos, sin):
	d = x.shape[-1]
	x1, x2 = x[..., : d // 2], x[..., d // 2:]
	cos = cos[None, None, :, :].to(x.dtype)
	sin = sin[None, None, :, :].to(x.dtype)
	rotated = torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
	return rotated


	class Attention(nn.Module):
	"""Multi-head self-attention with QK-Norm + RoPE + XSA."""

	def __init__(self, cfg):
	super().__init__()
	self.n_heads = cfg.n_heads
	self.head_dim = cfg.d_model // cfg.n_heads
	assert self.head_dim * cfg.n_heads == cfg.d_model

	self.wq = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
	self.wk = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
	self.wv = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
	self.wo = nn.Linear(cfg.d_model, cfg.d_model, bias=False)

	self.q_norm = RMSNorm(self.head_dim, eps=cfg.qknorm_eps)
	self.k_norm = RMSNorm(self.head_dim, eps=cfg.qknorm_eps)

	self.xsa_eps = 1e-6

	def forward(self, x, cos, sin):
	B, T, D = x.shape
	q = self.wq(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
	k = self.wk(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
	v = self.wv(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)

	q = self.q_norm(q)
	k = self.k_norm(k)

	q = apply_rope(q, cos[:T], sin[:T])
	k = apply_rope(k, cos[:T], sin[:T])

	y = F.scaled_dot_product_attention(q, k, v, is_causal=True)

	# --- XSA: z_i = y_i - (y_i . v_i)/(\|\|v_i\|\|^2+eps) * v_i, per token/head ---
	dot = (y * v).sum(dim=-1, keepdim=True)
	vnorm_sq = (v * v).sum(dim=-1, keepdim=True)
	proj_coef = dot / (vnorm_sq + self.xsa_eps)
	z = y - proj_coef * v

	z = z.transpose(1, 2).contiguous().view(B, T, D)
	return self.wo(z)


	class SwiGLU(nn.Module):
	def __init__(self, d_model, d_ffn):
	super().__init__()
	self.w_gate = nn.Linear(d_model, d_ffn, bias=False)
	self.w_up = nn.Linear(d_model, d_ffn, bias=False)
	self.w_down = nn.Linear(d_ffn, d_model, bias=False)

	def forward(self, x):
	return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))


	class Block(nn.Module):
	def __init__(self, cfg):
	super().__init__()
	self.attn_norm = RMSNorm(cfg.d_model, eps=cfg.rmsnorm_eps)
	self.attn = Attention(cfg)
	self.ffn_norm = RMSNorm(cfg.d_model, eps=cfg.rmsnorm_eps)
	self.ffn = SwiGLU(cfg.d_model, cfg.d_ffn)

	def forward(self, x, cos, sin):
	x = x + self.attn(self.attn_norm(x), cos, sin)
	x = x + self.ffn(self.ffn_norm(x))
	return x


	# ============================================================================
	# transformers-facing wrapper
	# ============================================================================

	class IvmeConversateModel(PreTrainedModel):
	"""Base model (no LM head) - exposed mainly so AutoModel works too;
	most users will want IvmeConversateForCausalLM below."""
	config_class = IvmeConversateConfig

	def __init__(self, config):
	super().__init__(config)
	self.tok_embed = nn.Embedding(config.vocab_size, config.d_model)
	self.blocks = nn.ModuleList([Block(config) for _ in range(config.n_layers)])
	self.final_norm = RMSNorm(config.d_model, eps=config.rmsnorm_eps)
	# NOTE: RoPE cos/sin are deliberately NOT cached as a
	# register_buffer here. transformers v5.x has a known, currently
	# open bug (confirmed via reproduction and matching upstream
	# reports: huggingface/transformers#43950, #44534) where
	# non-persistent buffers (persistent=False) get corrupted with
	# uninitialized/meta-device garbage memory during
	# from_pretrained()'s fast weight-materialization path - this
	# silently produced NaN/garbage RoPE values and completely wrong
	# model output when loaded via AutoModelForCausalLM, even though
	# the actual model weights loaded correctly. Recomputing cos/sin
	# fresh in forward() is cheap (a handful of trig ops) and
	# sidesteps this bug class entirely rather than depending on a
	# library fix.
	self._head_dim = config.d_model // config.n_heads
	self._rope_theta = config.rope_theta
	self._context_len = config.context_len

	self.post_init()

	def _init_weights(self, module):
	if isinstance(module, nn.Linear):
	nn.init.normal_(module.weight, mean=0.0, std=0.02)
	if module.bias is not None:
	nn.init.zeros_(module.bias)
	elif isinstance(module, nn.Embedding):
	nn.init.normal_(module.weight, mean=0.0, std=0.02)

	def forward(self, input_ids, **kwargs):
	# Defensive truncation: keep only the most recent context_len
	# tokens. This is a second safety layer beyond exposing
	# max_position_embeddings in the config (see
	# IvmeConversateConfig) - if any caller still manages to pass a
	# longer sequence, truncate cleanly to the trailing window
	# rather than crash inside RoPE with a cryptic tensor-shape
	# mismatch (confirmed reproducible: a 2048-token input against
	# this model's 1024 context previously crashed deep inside
	# apply_rope with no indication of the real cause).
	if input_ids.shape[1] > self._context_len:
	input_ids = input_ids[:, -self._context_len:]
	x = self.tok_embed(input_ids)
	cos, sin = precompute_rope(self._head_dim, self._context_len,
	theta=self._rope_theta, device=x.device)
	for block in self.blocks:
	x = block(x, cos, sin)
	x = self.final_norm(x)
	return x


	class IvmeConversateForCausalLM(PreTrainedModel, GenerationMixin):
	"""Standard transformers CausalLM interface: forward() returns
	CausalLMOutputWithPast(logits=...), supports labels for loss
	computation, and is compatible with model.generate() and any tool
	(like ArithMark's benchmark runner) expecting a normal
	AutoModelForCausalLM-loadable model."""

	config_class = IvmeConversateConfig
	# Tells transformers' save_pretrained/tie_weights machinery about the
	# tied embedding explicitly (required in transformers v5+; without
	# this, save_pretrained raises "shared tensors... not properly
	# defined" even though the weights ARE correctly tied via the
	# assignment in __init__ below - confirmed by reproducing the error
	# and fixing it with this attribute).
	_tied_weights_keys = {"lm_head.weight": "model.tok_embed.weight"}

	def __init__(self, config):
	super().__init__(config)
	self.model = IvmeConversateModel(config)
	self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)

	if config.tie_embeddings:
	self.lm_head.weight = self.model.tok_embed.weight

	self.post_init()

	def get_input_embeddings(self):
	return self.model.tok_embed

	def set_input_embeddings(self, value):
	self.model.tok_embed = value

	def get_output_embeddings(self):
	return self.lm_head

	def set_output_embeddings(self, new_embeddings):
	self.lm_head = new_embeddings

	def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
	# attention_mask accepted for interface compatibility but not
	# used - this model is trained/used causal-only with no padding
	# within a sequence (see ShardedTokenDataset / infer scripts).
	#
	# If input_ids exceeds context_len, self.model() truncates it
	# internally (see IvmeConversateModel.forward). labels must be
	# truncated identically here, BEFORE the loss is computed -
	# otherwise a caller passing both an over-length input_ids and
	# matching labels (e.g. a perplexity/loglikelihood task) would
	# hit a shape mismatch between logits and labels, or worse, a
	# silent misalignment if the shapes happened to coincide.
	context_len = self.config.max_position_embeddings
	if input_ids.shape[1] > context_len and labels is not None:
	labels = labels[:, -context_len:]

	hidden = self.model(input_ids)
	logits = self.lm_head(hidden)

	loss = None
	if labels is not None:
	loss = F.cross_entropy(
	logits.view(-1, logits.size(-1)), labels.view(-1), ignore_index=-100
	)

	return CausalLMOutputWithPast(
	loss=loss,
	logits=logits,
	)

	def prepare_inputs_for_generation(self, input_ids, **kwargs):
	# No KV-cache implementation - recomputes the full sequence each
	# step. Correct but not optimized for long generations; adequate
	# for benchmark scoring (log-likelihood over short continuations)
	# and interactive testing.
	return {"input_ids": input_ids}

	def generate(self, args, *kwargs):
	# Force use_cache=False unconditionally. This model has no cache
	# implementation, and transformers' generate() otherwise tries to
	# construct a DynamicCache using standard config field names
	# (num_hidden_layers, etc.) that this config doesn't define,
	# crashing regardless of the config-level use_cache default set
	# in IvmeConversateConfig (confirmed empirically: that default
	# does not reliably reach generate()'s cache-setup path in all
	# transformers versions). Overriding here guarantees correct
	# behavior for any caller - including external tools like
	# ArithMark's benchmark runner - with zero special knowledge
	# required on their end.
	kwargs["use_cache"] = False
	return super().generate(args, *kwargs)

	def num_params(self, non_embedding=False):
	n = sum(p.numel() for p in self.parameters())
	if non_embedding:
	n -= self.model.tok_embed.weight.numel()
	return n