File size: 871 Bytes
d53adc9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 | # Generated by ml.integrations.export.runtime_packager.write_remote_code_bundle.
# Exported for HuggingFace trust_remote_code loading.
# This file is intentionally self-contained.
from __future__ import annotations
from typing import Protocol
import torch
from torch import nn
class DecoderConfig(Protocol):
loss_chunk_size: int
return_logits_in_train: bool
vocab_size: int
class _OutputProjection(Protocol):
weight: torch.Tensor
class DecoderCoreModel(Protocol):
head_mixer: nn.Module
norm: nn.Module
output: _OutputProjection
gradient_checkpointing: bool
def _forward_hidden(
self,
input_ids: torch.Tensor,
start_pos: int = 0,
) -> tuple[torch.Tensor, torch.Tensor]: ...
def forward_full(self, input_ids: torch.Tensor) -> torch.Tensor: ...
__all__ = ["DecoderConfig", "DecoderCoreModel"]
|