Text Generation
Transformers
PyTorch
Safetensors
GGUF
English
collision
conversational
rag
reasoning
deepseek-r1-style
system-2
agent
ollama
llama.cpp
fastapi
openai-compatible
slm
edge-ai
cpu-first
in-house-nlp
math
keyphrase-extraction
topic-classification
grammar-correction
reading-comprehension
sentiment-analysis
research
educational
custom_code
| """ | |
| COLLISION Modeling for Hugging Face Transformers Integration. | |
| Enables AutoModelForCausalLM.from_pretrained("collision-10M/Collision-1B", trust_remote_code=True) | |
| and pipeline("text-generation", model="collision-10M/Collision-1B", trust_remote_code=True). | |
| """ | |
| import math | |
| from typing import Optional, Tuple, Union, List | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| try: | |
| from transformers.modeling_utils import PreTrainedModel | |
| from transformers.modeling_outputs import CausalLMOutputWithPast, BaseModelOutputWithPast | |
| from transformers.generation import GenerationMixin | |
| except ImportError: | |
| class PreTrainedModel(nn.Module): | |
| config_class = None | |
| base_model_prefix = "model" | |
| def __init__(self, config, *inputs, **kwargs): | |
| super().__init__() | |
| self.config = config | |
| class GenerationMixin: | |
| pass | |
| class CausalLMOutputWithPast: | |
| def __init__(self, loss=None, logits=None, past_key_values=None, hidden_states=None, attentions=None): | |
| self.loss = loss | |
| self.logits = logits | |
| self.past_key_values = past_key_values | |
| self.hidden_states = hidden_states | |
| self.attentions = attentions | |
| class BaseModelOutputWithPast: | |
| def __init__(self, last_hidden_state=None, past_key_values=None, hidden_states=None, attentions=None): | |
| self.last_hidden_state = last_hidden_state | |
| self.past_key_values = past_key_values | |
| self.hidden_states = hidden_states | |
| self.attentions = attentions | |
| try: | |
| from .configuration_collision import CollisionConfig | |
| except (ImportError, ValueError): | |
| from configuration_collision import CollisionConfig | |
| class TokenEmbedding(nn.Module): | |
| def __init__(self, vocab_size: int, d_model: int): | |
| super().__init__() | |
| self.emb = nn.Embedding(vocab_size, d_model) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| return self.emb(x) | |
| class PositionalEmbedding(nn.Module): | |
| def __init__(self, max_seq_len: int, d_model: int): | |
| super().__init__() | |
| self.emb = nn.Embedding(max_seq_len, d_model) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| seq_len = x.size(1) | |
| positions = torch.arange(seq_len, device=x.device).unsqueeze(0) | |
| return self.emb(positions) | |
| class CausalSelfAttention(nn.Module): | |
| def __init__(self, d_model: int, n_head: int, max_seq_len: int, dropout: float = 0.1): | |
| super().__init__() | |
| assert d_model % n_head == 0, "d_model must be divisible by n_head" | |
| self.d_model = d_model | |
| self.n_head = n_head | |
| self.d_k = d_model // n_head | |
| self.qkv_proj = nn.Linear(d_model, 3 * d_model) | |
| self.o_proj = nn.Linear(d_model, d_model) | |
| self.attn_dropout = nn.Dropout(dropout) | |
| self.resid_dropout = nn.Dropout(dropout) | |
| self.register_buffer( | |
| "bias", | |
| torch.tril(torch.ones(max_seq_len, max_seq_len)).view(1, 1, max_seq_len, max_seq_len), | |
| persistent=False | |
| ) | |
| def forward( | |
| self, | |
| x: torch.Tensor, | |
| attention_mask: Optional[torch.Tensor] = None, | |
| output_attentions: bool = False | |
| ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: | |
| B, T, C = x.size() | |
| q, k, v = self.qkv_proj(x).split(self.d_model, dim=2) | |
| q = q.view(B, T, self.n_head, self.d_k).transpose(1, 2) | |
| k = k.view(B, T, self.n_head, self.d_k).transpose(1, 2) | |
| v = v.view(B, T, self.n_head, self.d_k).transpose(1, 2) | |
| scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) | |
| scores = scores.masked_fill(self.bias[:, :, :T, :T] == 0, -1e4) | |
| if attention_mask is not None: | |
| if attention_mask.dim() == 2: | |
| # (B, T) -> (B, 1, 1, T) | |
| mask = attention_mask[:, None, None, :] | |
| scores = scores.masked_fill(mask == 0, -1e4) | |
| elif attention_mask.dim() == 4: | |
| scores = scores + attention_mask | |
| attn_weights = torch.softmax(scores, dim=-1) | |
| attn_weights = torch.nan_to_num(attn_weights, nan=0.0) | |
| attn_weights = self.attn_dropout(attn_weights) | |
| out = torch.matmul(attn_weights, v) | |
| out = out.transpose(1, 2).contiguous().view(B, T, C) | |
| out = self.o_proj(out) | |
| out = self.resid_dropout(out) | |
| outputs = (out,) | |
| if output_attentions: | |
| outputs = outputs + (attn_weights,) | |
| return outputs | |
| class FeedForward(nn.Module): | |
| def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1): | |
| super().__init__() | |
| self.net = nn.Sequential( | |
| nn.Linear(d_model, d_ff), | |
| nn.GELU(), | |
| nn.Linear(d_ff, d_model), | |
| nn.Dropout(dropout) | |
| ) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| return self.net(x) | |
| class TransformerBlock(nn.Module): | |
| def __init__(self, d_model: int, n_head: int, max_seq_len: int, d_ff: int, dropout: float = 0.1): | |
| super().__init__() | |
| self.ln1 = nn.LayerNorm(d_model) | |
| self.attn = CausalSelfAttention(d_model, n_head, max_seq_len, dropout) | |
| self.ln2 = nn.LayerNorm(d_model) | |
| self.ffn = FeedForward(d_model, d_ff, dropout) | |
| def forward( | |
| self, | |
| x: torch.Tensor, | |
| attention_mask: Optional[torch.Tensor] = None, | |
| output_attentions: bool = False | |
| ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: | |
| attn_outputs = self.attn(self.ln1(x), attention_mask=attention_mask, output_attentions=output_attentions) | |
| x = x + attn_outputs[0] | |
| x = x + self.ffn(self.ln2(x)) | |
| outputs = (x,) + attn_outputs[1:] | |
| return outputs | |
| class CollisionPreTrainedModel(PreTrainedModel): | |
| config_class = CollisionConfig | |
| base_model_prefix = "transformer" | |
| supports_gradient_checkpointing = True | |
| _no_split_modules = ["TransformerBlock"] | |
| def _init_weights(self, module): | |
| if isinstance(module, nn.Linear): | |
| torch.nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) | |
| if module.bias is not None: | |
| torch.nn.init.zeros_(module.bias) | |
| elif isinstance(module, nn.Embedding): | |
| emb_weight = module.emb.weight if hasattr(module, "emb") else module.weight | |
| torch.nn.init.normal_(emb_weight, mean=0.0, std=self.config.initializer_range) | |
| elif isinstance(module, nn.LayerNorm): | |
| torch.nn.init.ones_(module.weight) | |
| torch.nn.init.zeros_(module.bias) | |
| def post_init(self): | |
| self.apply(self._init_weights) | |
| class CollisionModel(CollisionPreTrainedModel): | |
| def __init__(self, config: CollisionConfig): | |
| super().__init__(config) | |
| self.config = config | |
| self.token_emb = TokenEmbedding(config.vocab_size, config.d_model) | |
| self.pos_emb = PositionalEmbedding(config.max_seq_len, config.d_model) | |
| self.drop = nn.Dropout(config.dropout) | |
| self.blocks = nn.ModuleList([ | |
| TransformerBlock( | |
| d_model=config.d_model, | |
| n_head=config.n_head, | |
| max_seq_len=config.max_seq_len, | |
| d_ff=config.d_ff, | |
| dropout=config.dropout | |
| ) for _ in range(config.n_layer) | |
| ]) | |
| self.ln_f = nn.LayerNorm(config.d_model, eps=config.layer_norm_epsilon) | |
| self.post_init() | |
| def get_input_embeddings(self): | |
| return self.token_emb.emb | |
| def set_input_embeddings(self, value): | |
| self.token_emb.emb = value | |
| def forward( | |
| self, | |
| input_ids: torch.LongTensor = None, | |
| attention_mask: Optional[torch.Tensor] = None, | |
| position_ids: Optional[torch.LongTensor] = None, | |
| inputs_embeds: Optional[torch.FloatTensor] = None, | |
| output_attentions: Optional[bool] = None, | |
| output_hidden_states: Optional[bool] = None, | |
| return_dict: Optional[bool] = None, | |
| **kwargs | |
| ) -> Union[Tuple, BaseModelOutputWithPast]: | |
| output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions if hasattr(self.config, "output_attentions") else False | |
| output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states if hasattr(self.config, "output_hidden_states") else False | |
| return_dict = return_dict if return_dict is not None else self.config.use_return_dict if hasattr(self.config, "use_return_dict") else True | |
| if input_ids is not None: | |
| batch_size, seq_length = input_ids.shape | |
| device = input_ids.device | |
| x = self.token_emb(input_ids) | |
| if position_ids is None: | |
| x = x + self.pos_emb(input_ids) | |
| else: | |
| x = x + self.pos_emb.emb(position_ids) | |
| elif inputs_embeds is not None: | |
| batch_size, seq_length, _ = inputs_embeds.shape | |
| device = inputs_embeds.device | |
| x = inputs_embeds | |
| if position_ids is None: | |
| positions = torch.arange(seq_length, device=device).unsqueeze(0) | |
| x = x + self.pos_emb.emb(positions) | |
| else: | |
| x = x + self.pos_emb.emb(position_ids) | |
| else: | |
| raise ValueError("You have to specify either input_ids or inputs_embeds") | |
| x = self.drop(x) | |
| all_hidden_states = () if output_hidden_states else None | |
| all_self_attns = () if output_attentions else None | |
| for block in self.blocks: | |
| if output_hidden_states: | |
| all_hidden_states = all_hidden_states + (x,) | |
| layer_outputs = block(x, attention_mask=attention_mask, output_attentions=output_attentions) | |
| x = layer_outputs[0] | |
| if output_attentions: | |
| all_self_attns = all_self_attns + (layer_outputs[1],) | |
| x = self.ln_f(x) | |
| if output_hidden_states: | |
| all_hidden_states = all_hidden_states + (x,) | |
| if not return_dict: | |
| return tuple(v for v in [x, all_hidden_states, all_self_attns] if v is not None) | |
| return BaseModelOutputWithPast( | |
| last_hidden_state=x, | |
| past_key_values=None, | |
| hidden_states=all_hidden_states, | |
| attentions=all_self_attns | |
| ) | |
| class CollisionForCausalLM(CollisionPreTrainedModel, GenerationMixin): | |
| _tied_weights_keys = ["lm_head.weight"] | |
| _keys_to_ignore_on_load_unexpected = {r"blocks\.\d+\.attn\.bias"} | |
| _keys_to_ignore_on_load_missing = set() | |
| def all_tied_weights_keys(self): | |
| return {"lm_head.weight": "token_emb.emb.weight"} | |
| def __init__(self, config: CollisionConfig): | |
| super().__init__(config) | |
| self.config = config | |
| # Direct embedding and transformer layers matching state_dict | |
| self.token_emb = TokenEmbedding(config.vocab_size, config.d_model) | |
| self.pos_emb = PositionalEmbedding(config.max_seq_len, config.d_model) | |
| self.drop = nn.Dropout(config.dropout) | |
| self.blocks = nn.ModuleList([ | |
| TransformerBlock( | |
| d_model=config.d_model, | |
| n_head=config.n_head, | |
| max_seq_len=config.max_seq_len, | |
| d_ff=config.d_ff, | |
| dropout=config.dropout | |
| ) for _ in range(config.n_layer) | |
| ]) | |
| self.ln_f = nn.LayerNorm(config.d_model, eps=config.layer_norm_epsilon) | |
| self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=True) | |
| if config.tie_embeddings: | |
| self.lm_head.weight = self.token_emb.emb.weight | |
| self.post_init() | |
| def get_input_embeddings(self): | |
| return self.token_emb.emb | |
| def set_input_embeddings(self, value): | |
| self.token_emb.emb = value | |
| def get_output_embeddings(self): | |
| return self.lm_head | |
| def set_output_embeddings(self, new_embeddings): | |
| self.lm_head = new_embeddings | |
| def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs): | |
| attention_mask = kwargs.get("attention_mask", None) | |
| position_ids = kwargs.get("position_ids", None) | |
| if input_ids.shape[1] > self.config.max_seq_len: | |
| input_ids = input_ids[:, -self.config.max_seq_len:] | |
| return { | |
| "input_ids": input_ids, | |
| "attention_mask": attention_mask, | |
| "position_ids": position_ids, | |
| } | |
| def forward( | |
| self, | |
| input_ids: torch.LongTensor = None, | |
| attention_mask: Optional[torch.Tensor] = None, | |
| position_ids: Optional[torch.LongTensor] = None, | |
| inputs_embeds: Optional[torch.FloatTensor] = None, | |
| labels: Optional[torch.LongTensor] = None, | |
| output_attentions: Optional[bool] = None, | |
| output_hidden_states: Optional[bool] = None, | |
| return_dict: Optional[bool] = None, | |
| **kwargs | |
| ) -> Union[Tuple, CausalLMOutputWithPast]: | |
| return_dict = return_dict if return_dict is not None else self.config.use_return_dict if hasattr(self.config, "use_return_dict") else True | |
| if input_ids is not None: | |
| B, T = input_ids.size() | |
| x = self.token_emb(input_ids) | |
| if position_ids is None: | |
| x = x + self.pos_emb(input_ids) | |
| else: | |
| x = x + self.pos_emb.emb(position_ids) | |
| elif inputs_embeds is not None: | |
| B, T, _ = inputs_embeds.size() | |
| x = inputs_embeds | |
| if position_ids is None: | |
| positions = torch.arange(T, device=inputs_embeds.device).unsqueeze(0) | |
| x = x + self.pos_emb.emb(positions) | |
| else: | |
| x = x + self.pos_emb.emb(position_ids) | |
| else: | |
| raise ValueError("You must specify either input_ids or inputs_embeds") | |
| x = self.drop(x) | |
| all_hidden_states = () if output_hidden_states else None | |
| all_self_attns = () if output_attentions else None | |
| for block in self.blocks: | |
| if output_hidden_states: | |
| all_hidden_states = all_hidden_states + (x,) | |
| layer_outputs = block(x, attention_mask=attention_mask, output_attentions=output_attentions) | |
| x = layer_outputs[0] | |
| if output_attentions: | |
| all_self_attns = all_self_attns + (layer_outputs[1],) | |
| hidden_states = self.ln_f(x) | |
| logits = self.lm_head(hidden_states) | |
| logits = torch.nan_to_num(logits, nan=0.0) | |
| loss = None | |
| if labels is not None: | |
| # Shift so that tokens < n predict n | |
| shift_logits = logits[..., :-1, :].contiguous() | |
| shift_labels = labels[..., 1:].contiguous() | |
| loss = F.cross_entropy( | |
| shift_logits.view(-1, shift_logits.size(-1)), | |
| shift_labels.view(-1), | |
| ignore_index=-100 | |
| ) | |
| if not return_dict: | |
| output = (logits,) + (all_hidden_states, all_self_attns) | |
| return ((loss,) + output) if loss is not None else output | |
| return CausalLMOutputWithPast( | |
| loss=loss, | |
| logits=logits, | |
| past_key_values=None, | |
| hidden_states=all_hidden_states, | |
| attentions=all_self_attns | |
| ) | |
| # Compatibility alias | |
| CollisionTransformer = CollisionForCausalLM | |