diff --git a/README.md b/README.md index 7cd85e0..6232095 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ a tiny vision language model that kicks ass and runs anywhere 1.6B parameter model built using SigLIP, Phi-1.5 and the LLaVA training dataset. Weights are licensed under CC-BY-SA due to using the LLaVA dataset. Try it out -on [Huggingface Spaces](https://huggingface.co/spaces/vikhyatk/moondream1)! +on [Hugging Face Spaces](https://huggingface.co/spaces/vikhyatk/moondream1)! **Benchmarks** diff --git a/moondream/__init__.py b/moondream/__init__.py index 85279b3..6662d86 100644 --- a/moondream/__init__.py +++ b/moondream/__init__.py @@ -1,2 +1,2 @@ -from .vision_encoder import VisionEncoder -from .text_model import TextModel +from .util import detect_device +from .moondream import Moondream diff --git a/moondream/phi/configuration_phi.py b/moondream/configuration_moondream.py similarity index 55% rename from moondream/phi/configuration_phi.py rename to moondream/configuration_moondream.py index b384c30..de16583 100644 --- a/moondream/phi/configuration_phi.py +++ b/moondream/configuration_moondream.py @@ -1,31 +1,20 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -import math -from typing import Optional - from transformers import PretrainedConfig +from typing import Optional +import math + class PhiConfig(PretrainedConfig): - """Phi configuration.""" - model_type = "phi-msft" - attribute_map = { - "max_position_embeddings": "n_positions", - "hidden_size": "n_embd", - "num_attention_heads": "n_head", - "num_hidden_layers": "n_layer", - } def __init__( self, - vocab_size: int = 50304, + vocab_size: int = 51200, n_positions: int = 2048, - n_embd: int = 1024, - n_layer: int = 20, + n_embd: int = 2048, + n_layer: int = 24, n_inner: Optional[int] = None, - n_head: int = 16, + n_head: int = 32, n_head_kv: Optional[int] = None, rotary_dim: Optional[int] = 32, activation_function: Optional[str] = "gelu_new", @@ -41,26 +30,45 @@ class PhiConfig(PretrainedConfig): pad_vocab_size_multiple: int = 64, gradient_checkpointing: bool = False, **kwargs - ) -> None: - self.vocab_size = int( + ): + pad_vocab_size = ( math.ceil(vocab_size / pad_vocab_size_multiple) * pad_vocab_size_multiple ) - self.n_positions = n_positions - self.n_embd = n_embd - self.n_layer = n_layer - self.n_inner = n_inner - self.n_head = n_head - self.n_head_kv = n_head_kv + super().__init__( + vocab_size=pad_vocab_size, + n_positions=n_positions, + n_embd=n_embd, + n_layer=n_layer, + n_inner=n_inner, + n_head=n_head, + n_head_kv=n_head_kv, + activation_function=activation_function, + attn_pdrop=attn_pdrop, + embd_pdrop=embd_pdrop, + resid_pdrop=resid_pdrop, + layer_norm_epsilon=layer_norm_epsilon, + initializer_range=initializer_range, + pad_vocab_size_multiple=pad_vocab_size_multiple, + tie_word_embeddings=tie_word_embeddings, + gradient_checkpointing=gradient_checkpointing, + **kwargs + ) self.rotary_dim = min(rotary_dim, n_embd // n_head) - self.activation_function = activation_function self.flash_attn = flash_attn self.flash_rotary = flash_rotary self.fused_dense = fused_dense - self.attn_pdrop = attn_pdrop - self.embd_pdrop = embd_pdrop - self.resid_pdrop = resid_pdrop - self.layer_norm_epsilon = layer_norm_epsilon - self.initializer_range = initializer_range - self.gradient_checkpointing = gradient_checkpointing - super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) + attribute_map = { + "max_position_embeddings": "n_positions", + "hidden_size": "n_embd", + "num_attention_heads": "n_head", + "num_hidden_layers": "n_layer", + } + + +class MoondreamConfig(PretrainedConfig): + model_type = "moondream1" + + def __init__(self, **kwargs): + self.phi_config = PhiConfig(**kwargs) + super().__init__(**kwargs) diff --git a/moondream/modeling_phi.py b/moondream/modeling_phi.py new file mode 100644 index 0000000..ef54625 --- /dev/null +++ b/moondream/modeling_phi.py @@ -0,0 +1,720 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +# +# Copyright (c) 2022, Tri Dao, trid@cs.stanford.edu. +# Licensed under the BSD 3-Clause License. + +from dataclasses import dataclass, field +from typing import Any, Dict, Optional, Union, Tuple + +import math +import torch +import torch.nn as nn +from einops import rearrange, repeat +from transformers import PretrainedConfig, PreTrainedModel +from transformers.activations import ACT2FN +from transformers.modeling_outputs import CausalLMOutputWithPast + +from .configuration_moondream import PhiConfig + +FusedDense = None + + +@dataclass +class InferenceParams: + max_seqlen: int + max_batch_size: int + seqlen_offset: int = 0 + batch_size_offset: int = 0 + key_value_memory_dict: Dict[str, Any] = field(default_factory=dict) + lengths_per_sample: torch.Tensor = None + + +class Embedding(nn.Module): + def __init__(self, config: PretrainedConfig): + super().__init__() + self.wte = nn.Embedding(config.vocab_size, config.n_embd) + self.drop = nn.Dropout(config.embd_pdrop) + + def forward(self, input_ids: torch.LongTensor) -> torch.FloatTensor: + return self.drop(self.wte(input_ids.view(-1, input_ids.size(-1)))) + + +def _apply_rotary_emb(x, cos, sin): + seqlen, rotary_dim = x.size(1), cos.size(1) * 2 + x_rot, x_pass = x[..., :rotary_dim], x[..., rotary_dim:] + x1, x2 = x_rot.chunk(2, dim=-1) + c, s = cos[:seqlen].unsqueeze(1), sin[:seqlen].unsqueeze(1) + x_rot = torch.cat([x1 * c - x2 * s, x1 * s + x2 * c], dim=-1) + return torch.cat([x_rot.to(x.dtype), x_pass], dim=-1) + + +def _apply_rotary_emb_kv( + kv: torch.FloatTensor, cos: torch.FloatTensor, sin: torch.FloatTensor +) -> torch.FloatTensor: + seqlen, rotary_dim = kv.shape[1], cos.shape[-1] * 2 + k_rot = kv[:, :, 0, :, :rotary_dim].chunk(2, dim=-1) + k_pass = kv[:, :, 0, :, rotary_dim:] + c, s = cos[:seqlen].unsqueeze(1), sin[:seqlen].unsqueeze(1) + k_rot = torch.cat( + [k_rot[0] * c - k_rot[1] * s, k_rot[0] * s + k_rot[1] * c], dim=-1 + ) + return torch.cat( + [torch.cat([k_rot, k_pass], dim=-1).unsqueeze(2), kv[:, :, 1:2, :, :]], dim=2 + ) + + +def _apply_rotary_emb_qkv( + qkv: torch.FloatTensor, cos: torch.FloatTensor, sin: torch.FloatTensor +) -> torch.FloatTensor: + seqlen, rotary_dim = qkv.shape[1], cos.shape[1] * 2 + + c = cos[:seqlen].unsqueeze(1) + s = sin[:seqlen].unsqueeze(1) + + qkv_rot = torch.stack( + [ + torch.cat( + [ + qkv[:, :, i, :, : rotary_dim // 2] * c + - qkv[:, :, i, :, rotary_dim // 2 : rotary_dim] * s, + qkv[:, :, i, :, : rotary_dim // 2] * s + + qkv[:, :, i, :, rotary_dim // 2 : rotary_dim] * c, + ], + dim=-1, + ).to(qkv.dtype) + for i in range(2) + ], + dim=2, + ) + + qkv_pass = qkv[:, :, :2, :, rotary_dim:].unsqueeze(2) + qkv_v = qkv[:, :, 2:3, :, :] + return torch.cat([qkv_rot, qkv_pass, qkv_v], dim=2) + + +class RotaryEmbedding(nn.Module): + # Enhanced Transformer with Rotary Position Embedding (https://arxiv.org/pdf/2104.09864.pdf) + def __init__( + self, + dim: int, + base: int = 10000, + scale_base: Optional[float] = None, + pos_idx_in_fp32: bool = True, + max_position_embeddings: int = 2048, + device: Optional[str] = None, + ) -> None: + super().__init__() + # fp32 is preferred since the output of `torch.arange` can be quite large and bf16 would lose a lot of precision + self.dim, self.base, self.pos_idx_in_fp32, self.device = ( + dim, + float(base), + pos_idx_in_fp32, + device, + ) + self.max_position_embeddings = max_position_embeddings + if scale_base is not None: + raise NotImplementedError + + # Generate and register the non-trainable buffers + self.register_buffer( + "inv_freq", self._compute_inv_freq(device), persistent=False + ) + self.register_buffer( + "scale", self._calculate_scale(dim, scale_base, device), persistent=False + ) + self._update_cos_sin_cache( + max_position_embeddings, device=device, dtype=torch.float32 + ) + + def _calculate_scale(self, dim, scale_base, device): + return ( + ( + ( + torch.arange(0, dim, 2, device=device, dtype=torch.float32) + + 0.4 * dim + ) + / (1.4 * dim) + ) + if scale_base is not None + else None + ) + + def _compute_inv_freq(self, device: Optional[str] = None) -> torch.FloatTensor: + return 1.0 / ( + self.base + ** ( + torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) + / self.dim + ) + ) + + def _update_cos_sin_cache( + self, + seqlen: int, + device: Optional[str] = None, + dtype: Optional[torch.dtype] = None, + ) -> None: + self._seq_len_cached = seqlen + t = torch.arange( + seqlen, + device=device, + dtype=torch.float32 if self.pos_idx_in_fp32 else self.inv_freq.dtype, + ) + inv_freq = ( + self._compute_inv_freq(device=device) + if self.pos_idx_in_fp32 and self.inv_freq.dtype != torch.float32 + else self.inv_freq + ) + + freqs = torch.outer(t, inv_freq) + + def apply_scale(freqs, scale, operator, dtype): + result = operator(freqs) + return (result / scale).to(dtype) if scale is not None else result.to(dtype) + + if scale := self.scale: + power = ( + torch.arange(seqlen, dtype=scale.dtype, device=scale.device) + - seqlen // 2 + ) / self.scale_base + scale = scale.to(device=power.device) ** power.unsqueeze(1) + + self._cos_cached = apply_scale( + freqs, 1 / scale if scale is not None else None, torch.cos, dtype + ) + self._sin_cached = apply_scale( + freqs, 1 / scale if scale is not None else None, torch.sin, dtype + ) + if scale is not None: + self._cos_k_cached = apply_scale(freqs, scale, torch.cos, dtype) + self._sin_k_cached = apply_scale(freqs, scale, torch.sin, dtype) + + def forward( + self, + qkv: torch.Tensor, + kv: Optional[torch.Tensor] = None, + seqlen_offset: int = 0, + ) -> Tuple[torch.Tensor, torch.Tensor]: + should_update = ( + self._seq_len_cached < qkv.shape[1] + seqlen_offset + or self._cos_cached.device != qkv.device + or self._cos_cached.dtype != qkv.dtype + or (self.training and self._cos_cached.is_inference()) + ) + + if should_update: + self._update_cos_sin_cache( + qkv.shape[1] + seqlen_offset, device=qkv.device, dtype=qkv.dtype + ) + + offset_cos = self._cos_cached[seqlen_offset:] + offset_sin = self._sin_cached[seqlen_offset:] + + if kv is None: + return _apply_rotary_emb_qkv(qkv, offset_cos, offset_sin) + else: + return _apply_rotary_emb(qkv, offset_cos, offset_sin), _apply_rotary_emb_kv( + kv, offset_cos, offset_sin + ) + + +class MLP(nn.Module): + def __init__( + self, + config: PretrainedConfig, + n_inner: Optional[int] = None, + act_fn: Optional[str] = None, + ) -> None: + super().__init__() + n_inner = n_inner or getattr(config, "n_inner", None) or 4 * config.n_embd + act_fn = act_fn or config.activation_function + + self.fc1 = nn.Linear(config.n_embd, n_inner) + self.fc2 = nn.Linear(n_inner, config.n_embd) + self.act = ACT2FN[act_fn] + + def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor: + return self.fc2(self.act(self.fc1(hidden_states))) + + +# Flash Attention (https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py) +class SelfAttention(nn.Module): + def __init__( + self, + causal: bool = True, + softmax_scale: Optional[float] = None, + attention_dropout: float = 0.0, + ): + super().__init__() + self.causal = causal + self.softmax_scale = softmax_scale + self.drop = nn.Dropout(attention_dropout) + + @torch.autocast("cpu", enabled=False) + @torch.autocast("cuda", enabled=False) + def forward( + self, + qkv: torch.FloatTensor, + causal: Optional[bool] = None, + key_padding_mask: Optional[torch.BoolTensor] = None, + ): + q, k, v = qkv.chunk(3, dim=-1) + scale = self.softmax_scale or 1.0 / q.size(-1) ** 0.5 + + scores = ( + torch.einsum("bthd,bshd->bhts", q.to(torch.float32), k.to(torch.float32)) + * scale + ) + if causal or self.causal: + scores.triu_(1).fill_(-10000.0) + if key_padding_mask is not None: + scores.masked_fill_(key_padding_mask[:, None, None, :], -10000.0) + + attn = self.drop(torch.softmax(scores, dim=-1).to(v.dtype)) + return torch.einsum("bhts,bshd->bthd", attn, v) + + +# Flash Attention (https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py) +class CrossAttention(nn.Module): + def __init__(self, causal=True, softmax_scale=None, attention_dropout=0.0): + super().__init__() + self.causal = causal + self.softmax_scale = softmax_scale + self.drop = nn.Dropout(attention_dropout) + + @torch.autocast("cpu", enabled=False) + @torch.autocast("cuda", enabled=False) + def forward( + self, + q: torch.FloatTensor, + kv: torch.FloatTensor, + causal: bool = None, + key_padding_mask: Optional[torch.BoolTensor] = None, + ) -> torch.FloatTensor: + batch_size, seqlen_q = q.shape[0], q.shape[1] + seqlen_k = kv.shape[1] + + if kv.shape[3] != q.shape[2]: + kv = repeat(kv, "... hkv d -> ... (hkv g) d", g=q.shape[2] // kv.shape[3]) + k, v = kv.unbind(dim=2) + + q = q.to(torch.float32) + k = k.to(torch.float32) + + causal = self.causal if causal is None else causal + softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1]) + + # Autocast is manually disabled to avoid `torch.einsum` performing the operation using float16, which might lead to overflow + scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale) + + if key_padding_mask is not None: + padding_mask = torch.full( + (batch_size, seqlen_k), + -10000.0, + dtype=scores.dtype, + device=scores.device, + ) + padding_mask.masked_fill_(key_padding_mask, 0.0) + scores = scores + rearrange(padding_mask, "b s -> b 1 1 s") + + if causal: + rows = rearrange( + torch.arange(seqlen_q, device=q.device, dtype=torch.long), "s -> s 1" + ) + cols = torch.arange(seqlen_k, device=k.device, dtype=torch.long) + causal_mask = cols > rows + seqlen_k - seqlen_q + scores = scores.masked_fill(causal_mask, -10000.0) + + attention = torch.softmax(scores, dim=-1).to(v.dtype) + attention = self.drop(attention) + output = torch.einsum("bhts,bshd->bthd", attention, v) + + return output + + +def _find_mha_dims( + config: PretrainedConfig, + n_head: Optional[int] = None, + n_head_kv: Optional[int] = None, + head_dim: Optional[int] = None, +) -> Tuple[int, int]: + if n_head is None and head_dim is None: + head_dim = config.n_embd // config.n_head + n_head = config.n_head + elif n_head is None or head_dim is None: + raise ValueError("`n_head` and `head_dim` must be both specified or `None`.") + if n_head_kv is None: + n_head_kv = getattr(config, "n_head_kv", None) or n_head + return n_head, n_head_kv, head_dim + + +def _update_kv_cache( + kv: torch.FloatTensor, inference_params: InferenceParams, layer_idx: int +) -> torch.FloatTensor: + num_heads, head_dim = kv.shape[-2:] + layer_memory = inference_params.key_value_memory_dict.setdefault( + layer_idx, + torch.empty( + inference_params.max_batch_size, + inference_params.max_seqlen, + 2, + num_heads, + head_dim, + dtype=kv.dtype, + device=kv.device, + ), + ) + + batch_slice = slice( + inference_params.batch_size_offset, + inference_params.batch_size_offset + kv.shape[0], + ) + seqlen_slice = slice( + inference_params.seqlen_offset, inference_params.seqlen_offset + kv.shape[1] + ) + + if seqlen_slice.stop >= inference_params.max_seqlen: + layer_memory = torch.cat((layer_memory, kv), dim=1) + inference_params.key_value_memory_dict[layer_idx] = layer_memory + + layer_memory[batch_slice, seqlen_slice, ...] = kv + return layer_memory[batch_slice, : seqlen_slice.stop, ...] + + +# Multi-head attention layer with rotary embeddings +class MHA(nn.Module): + def __init__( + self, + config, + dtype=None, + device=None, + rotary_dim=None, + rotary_base=10000.0, + rotary_scale_base=None, + n_head=None, + n_head_kv=None, + head_dim=None, + bias=True, + causal=True, + softmax_scale=None, + layer_idx=None, + return_residual=False, + checkpointing=False, + ): + super().__init__() + + # Set rotary embedding if specified + self.rotary_dim = rotary_dim or getattr(config, "rotary_dim", 0) + if self.rotary_dim: + self.rotary_emb = RotaryEmbedding( + self.rotary_dim, + base=rotary_base, + scale_base=rotary_scale_base, + device=device, + max_position_embeddings=config.n_positions, + ) + + # Determine MHA dims from arguments or config + self.n_head, self.n_head_kv, self.head_dim = _find_mha_dims( + config, n_head, n_head_kv, head_dim + ) + op_size = self.head_dim * (self.n_head + 2 * self.n_head_kv) + hidden_size = config.n_embd + + # Choose Linear class based on config, FusedDense is optional + LinearClass = ( + FusedDense if config.fused_dense and FusedDense is not None else nn.Linear + ) + self.Wqkv = LinearClass( + hidden_size, op_size, bias=bias, device=device, dtype=dtype + ) + self.out_proj = LinearClass( + hidden_size, hidden_size, bias=bias, device=device, dtype=dtype + ) + + # Initialize attention mechanisms + attn_kwargs = { + "causal": causal, + "softmax_scale": softmax_scale, + "attention_dropout": config.attn_pdrop, + } + self.inner_attn = SelfAttention(**attn_kwargs) + self.inner_cross_attn = CrossAttention(**attn_kwargs) + + self.layer_idx = layer_idx + self.return_residual = return_residual + self.checkpointing = checkpointing + + def _forward_self_attn( + self, x: torch.FloatTensor, key_padding_mask: Optional[torch.BoolTensor] + ) -> torch.FloatTensor: + qkv = rearrange( + self.Wqkv(x), "... (three h d) -> ... three h d", three=3, d=self.head_dim + ) + if self.rotary_dim > 0: + qkv = self.rotary_emb(qkv) + attn_func = ( + torch.utils.checkpoint.checkpoint + if self.checkpointing + else lambda f, *args, **kwargs: f(*args, **kwargs) + ) + return attn_func(self.inner_attn, qkv, key_padding_mask=key_padding_mask) + + def _forward_cross_attn( + self, + x: torch.FloatTensor, + past_key_values: Optional[InferenceParams], + key_padding_mask: Optional[torch.BoolTensor], + ) -> torch.FloatTensor: + qkv = self.Wqkv(x) + q, kv = ( + qkv[..., : self.n_head * self.head_dim], + qkv[..., self.n_head * self.head_dim :], + ) + q = rearrange(q, "... (h d) -> ... h d", d=self.head_dim) + kv = rearrange(kv, "... (two hkv d) -> ... two hkv d", two=2, d=self.head_dim) + + seqlen_offset = ( + past_key_values.seqlen_offset if past_key_values is not None else 0 + ) + causal = None if seqlen_offset == 0 else False + if self.rotary_dim > 0: + q, kv = self.rotary_emb(q, kv=kv, seqlen_offset=seqlen_offset) + + if past_key_values is not None: + kv = _update_kv_cache(kv, past_key_values, self.layer_idx) + + attn_func = ( + torch.utils.checkpoint.checkpoint + if self.checkpointing + else lambda fn, *args, **kwargs: fn(*args, **kwargs) + ) + + return attn_func( + self.inner_cross_attn, + q, + kv, + key_padding_mask=key_padding_mask, + causal=causal, + ) + + def forward( + self, + x: torch.FloatTensor, + past_key_values: Optional[InferenceParams] = None, + attention_mask: Optional[Union[torch.LongTensor, torch.BoolTensor]] = None, + ) -> Tuple[torch.FloatTensor, torch.FloatTensor]: + attention_mask = attention_mask.bool() if attention_mask is not None else None + use_cross_attn = self.n_head != self.n_head_kv or past_key_values is not None + attn_output_function = ( + self._forward_cross_attn if use_cross_attn else self._forward_self_attn + ) + attn_output = ( + attn_output_function(x, past_key_values, attention_mask) + if use_cross_attn + else attn_output_function(x, attention_mask) + ) + output = self.out_proj(rearrange(attn_output, "... h d -> ... (h d)")) + return (output, x) if self.return_residual else output + + +# Parallel block. This block applies parallel mixer and MLP layers to the input (used in GPT-J and CodeGen). +class ParallelBlock(nn.Module): + def __init__(self, config: PretrainedConfig, block_idx: Optional[int] = None): + super().__init__() + self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon) + self.resid_dropout = nn.Dropout(config.resid_pdrop) + self.block_idx = block_idx + self.mixer = MHA(config, layer_idx=block_idx) + self.mlp = MLP(config) + + def forward( + self, + hidden_states: torch.FloatTensor, + past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None, + attention_mask: Optional[torch.BoolTensor] = None, + ) -> torch.FloatTensor: + residual = hidden_states + hidden_states = self.ln(hidden_states) + + attn_outputs = self.mixer( + hidden_states, + past_key_values=past_key_values, + attention_mask=attention_mask, + ) + if isinstance(attn_outputs, tuple): + attn_outputs = attn_outputs[0] + + attn_outputs = self.resid_dropout(attn_outputs) + feed_forward_hidden_states = self.resid_dropout(self.mlp(hidden_states)) + return attn_outputs + feed_forward_hidden_states + residual + + +class CausalLMHead(nn.Module): + """Causal Language Modeling head. Simplified version.""" + + def __init__(self, config): + super().__init__() + self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon) + self.linear = nn.Linear(config.n_embd, config.vocab_size) + + def forward(self, hidden_states): + return self.linear(self.ln(hidden_states)).to(torch.float32) + + +# Improving Language Understanding by Generative Pre-Training +# (https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf) +class CausalLMLoss(nn.Module): + def __init__(self, shift_labels: bool = True) -> None: + super().__init__() + self.shift_labels = shift_labels + self.loss_fct = nn.CrossEntropyLoss() + + def forward( + self, logits: torch.FloatTensor, labels: torch.LongTensor + ) -> torch.FloatTensor: + if self.shift_labels: + logits, labels = logits[..., :-1, :], labels[..., 1:] + return self.loss_fct(logits.reshape(-1, logits.size(-1)), labels.reshape(-1)) + + +class PhiPreTrainedModel(PreTrainedModel): + config_class = PhiConfig + base_model_prefix = "transformer" + supports_gradient_checkpointing = False + _no_split_modules = ["ParallelBlock"] + + def __init__(self, *inputs, **kwargs) -> None: + super().__init__(*inputs, **kwargs) + + def prepare_inputs_for_generation( + self, + input_ids: torch.LongTensor = None, + inputs_embeds: torch.FloatTensor = None, + past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None, + attention_mask: Optional[Union[torch.LongTensor, torch.BoolTensor]] = None, + **kwargs, + ) -> Dict[str, Any]: + if input_ids is None and inputs_embeds is None: + raise ValueError( + "You have to specify either `input_ids` or `inputs_embeds`." + ) + + max_batch_size = ( + inputs_embeds.shape[0] if inputs_embeds is not None else input_ids.shape[0] + ) + seqlen_offset = ( + inputs_embeds.shape[1] + input_ids.shape[1] - 2 + if inputs_embeds is not None + else input_ids.shape[1] - 1 + ) + + args = ( + {"inputs_embeds": inputs_embeds} + if inputs_embeds is not None + else {"input_ids": input_ids} + ) + + if not isinstance(past_key_values, InferenceParams): + past_key_values = InferenceParams( + max_seqlen=self.config.n_positions, + max_batch_size=max_batch_size, + seqlen_offset=0, + batch_size_offset=0, + key_value_memory_dict={}, + lengths_per_sample=None, + ) + else: + past_key_values.seqlen_offset = seqlen_offset + args = {"input_ids": input_ids[:, -1].unsqueeze(-1)} + + return { + **args, + "past_key_values": past_key_values, + "attention_mask": attention_mask, + } + + +class PhiModel(PhiPreTrainedModel): + _keys_to_ignore_on_load_missing = [""] + _keys_to_ignore_on_load_unexpected = [r"h\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)"] + + def __init__(self, config: PhiConfig) -> None: + super().__init__(config) + self.embd = Embedding(config) + self.h = nn.ModuleList( + [ParallelBlock(config, block_idx=i) for i in range(config.n_layer)] + ) + self.gradient_checkpointing = config.gradient_checkpointing + self.post_init() + + def get_input_embeddings(self) -> nn.Embedding: + return self.embd.wte + + def set_input_embeddings(self, new_embeddings: nn.Embedding) -> None: + self.embd.wte = new_embeddings + + def forward( + self, + input_ids: torch.LongTensor = None, + inputs_embeds: torch.FloatTensor = None, + past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None, + attention_mask: Optional[torch.BoolTensor] = None, + ) -> torch.FloatTensor: + if (input_ids is None) == (inputs_embeds is None): + raise ValueError("Specify exactly one of `input_ids` or `inputs_embeds`.") + hidden_states = self.embd(input_ids) if input_ids is not None else inputs_embeds + + for layer in self.h: + func = layer.__call__ if self.gradient_checkpointing else layer + args = (hidden_states, past_key_values, attention_mask) + hidden_states = ( + torch.utils.checkpoint.checkpoint(func, *args, use_reentrant=True) + if self.gradient_checkpointing + else func(*args) + ) + + return hidden_states + + +class PhiForCausalLM(PhiPreTrainedModel): + _keys_to_ignore_on_load_missing, _keys_to_ignore_on_load_unexpected = ( + [""], + [r"transformer\.h\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)"], + ) + + def __init__(self, config: PhiConfig) -> None: + super().__init__(config) + self.transformer = PhiModel(config) + self.lm_head = CausalLMHead(config) + self.loss = CausalLMLoss() + self.post_init() + + def get_output_embeddings(self) -> nn.Linear: + return self.lm_head.linear + + def set_output_embeddings(self, new_embeddings: nn.Linear) -> None: + self.lm_head.linear = new_embeddings + + def forward( + self, + input_ids: torch.LongTensor = None, + inputs_embeds: torch.FloatTensor = None, + past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None, + attention_mask: Optional[torch.BoolTensor] = None, + labels: Optional[torch.LongTensor] = None, + **kwargs, + ) -> CausalLMOutputWithPast: + hidden_states = self.transformer( + input_ids=input_ids, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + attention_mask=attention_mask, + ) + lm_logits = self.lm_head(hidden_states) + loss = self.loss(lm_logits, labels) if labels is not None else None + + return CausalLMOutputWithPast( + loss=loss, logits=lm_logits, past_key_values=past_key_values + ) diff --git a/moondream/moondream.py b/moondream/moondream.py new file mode 100644 index 0000000..9a223d4 --- /dev/null +++ b/moondream/moondream.py @@ -0,0 +1,100 @@ +import torch +from .vision_encoder import VisionEncoder +from .text_model import TextModel +from .configuration_moondream import MoondreamConfig +from transformers import PreTrainedModel +import re + + +class Moondream(PreTrainedModel): + config_class = MoondreamConfig + + def __init__(self, config): + super().__init__(config) + self.vision_encoder = VisionEncoder() + self.text_model = TextModel(config) + + @property + def device(self): + return self.text_model.model.device + + def encode_image(self, image): + return self.vision_encoder(image) + + def input_embeds(self, prompt, image_embeds, tokenizer): + def _tokenize(txt): + return tokenizer( + txt, return_tensors="pt", add_special_tokens=False + ).input_ids.to(self.device) + + # Add BOS token + embeds = [] + embeds.append( + self.text_model.text_emb( + (torch.tensor([[tokenizer.bos_token_id]], device=self.device)) + ) + ) + + if "" not in prompt: + embeds.append(self.text_model.text_emb(_tokenize(prompt))) + else: + assert prompt.count("") == 1 + before, after = prompt.split("") + embeds.append(self.text_model.text_emb(_tokenize(f"{before}"))) + embeds.append(image_embeds.to(self.device)) + embeds.append(self.text_model.text_emb(_tokenize(f"{after}"))) + + return torch.cat(embeds, dim=1) + + def generate( + self, + image_embeds, + prompt, + tokenizer, + eos_text="Human:", + max_new_tokens=128, + **kwargs, + ): + eos_tokens = tokenizer(eos_text, add_special_tokens=False)[0].ids + + generate_config = { + "eos_token_id": eos_tokens, + "bos_token_id": tokenizer.bos_token_id, + "pad_token_id": tokenizer.eos_token_id, + "max_new_tokens": max_new_tokens, + **kwargs, + } + + with torch.no_grad(): + inputs_embeds = self.input_embeds(prompt, image_embeds, tokenizer) + output_ids = self.text_model.model.generate( + inputs_embeds=inputs_embeds, **generate_config + ) + + return tokenizer.batch_decode(output_ids, skip_special_tokens=True) + + def answer_question( + self, + image_embeds, + question, + tokenizer, + chat_history="", + result_queue=None, + **kwargs, + ): + prompt = f"\n\n{chat_history}Question: {question}\n\nAnswer:" + answer = self.generate( + image_embeds, + prompt, + eos_text="", + tokenizer=tokenizer, + max_new_tokens=128, + **kwargs, + )[0] + cleaned_answer = re.sub("<$", "", re.sub("END$", "", answer)).strip() + + # Use the result_queue to pass the result if it is provided + if result_queue: + result_queue.put(cleaned_answer) + else: + return cleaned_answer diff --git a/moondream/phi/modeling_phi.py b/moondream/phi/modeling_phi.py deleted file mode 100644 index b7317f4..0000000 --- a/moondream/phi/modeling_phi.py +++ /dev/null @@ -1,970 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. -# -# Copyright (c) 2022, Tri Dao, trid@cs.stanford.edu. -# Licensed under the BSD 3-Clause License. - -from __future__ import annotations - -import math -from dataclasses import dataclass, field -from typing import Any, Dict, Optional, Tuple, Union - -import torch -import torch.nn as nn -from einops import rearrange, repeat -from transformers import PretrainedConfig, PreTrainedModel -from transformers.activations import ACT2FN -from transformers.modeling_outputs import CausalLMOutputWithPast - -from .configuration_phi import PhiConfig - -pad_input, unpad_input = None, None -FlashRotaryEmbedding = None -FlashSelfAttention, FlashCrossAttention = None, None -FusedDense = None - - -@dataclass -class InferenceParams: - """Inference parameters passed to model to efficiently calculate - and store context during inference. - - Reference: - https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/utils/generation.py. - - Args: - max_seqlen: Maximum sequence length. - max_batch_size: Maximum batch size. - seqlen_offset: Sequence length offset. - batch_size_offset: Batch size offset. - key_value_memory_dict: Key value memory dictionary. - lengths_per_sample: Lengths per sample. - - """ - - max_seqlen: int = field(metadata={"help": "Maximum sequence length."}) - - max_batch_size: int = field(metadata={"help": "Maximum batch size."}) - - seqlen_offset: int = field(default=0, metadata={"help": "Sequence length offset."}) - - batch_size_offset: int = field(default=0, metadata={"help": "Batch size offset."}) - - key_value_memory_dict: Dict[str, Any] = field( - default_factory=dict, metadata={"help": "Key value memory dictionary."} - ) - - lengths_per_sample: torch.Tensor = field( - default=None, metadata={"help": "Lengths per sample."} - ) - - -class Embedding(nn.Module): - """Token embedding with dropout.""" - - def __init__(self, config: PretrainedConfig) -> None: - super().__init__() - - self.wte = nn.Embedding(config.vocab_size, config.n_embd) - self.drop = nn.Dropout(config.embd_pdrop) - - def forward(self, input_ids: torch.LongTensor) -> torch.FloatTensor: - input_shape = input_ids.size() - input_ids = input_ids.view(-1, input_shape[-1]) - - hidden_states = self.wte(input_ids) - hidden_states = self.drop(hidden_states) - - return hidden_states - - -# @torch.compile -def _apply_rotary_emb( - x: torch.FloatTensor, - cos: torch.FloatTensor, - sin: torch.FloatTensor, -) -> torch.FloatTensor: - _, seqlen, _, _ = x.shape - _, rotary_dim = cos.shape - rotary_dim *= 2 - - x_rot = x[:, :, :, :rotary_dim] - x_pass = x[:, :, :, rotary_dim:] - - x1, x2 = x_rot.chunk(2, dim=-1) - c, s = rearrange(cos[:seqlen], "s d -> s 1 d"), rearrange( - sin[:seqlen], "s d -> s 1 d" - ) - x1, x2, c, s = [t.to(dtype=torch.float32) for t in [x1, x2, c, s]] - - x_rot = torch.cat([x1 * c - x2 * s, x1 * s + x2 * c], axis=-1).to(x.dtype) - - return torch.cat([x_rot, x_pass], axis=-1) - - -# @torch.compile -def _apply_rotary_emb_kv( - kv: torch.FloatTensor, - cos: torch.FloatTensor, - sin: torch.FloatTensor, - cos_k: Optional[torch.FloatTensor] = None, - sin_k: Optional[torch.FloatTensor] = None, -) -> torch.FloatTensor: - _, seqlen, _, _, _ = kv.shape - _, rotary_dim = cos.shape - rotary_dim *= 2 - - k_rot = kv[:, :, 0, :, :rotary_dim] - k_pass = kv[:, :, 0, :, rotary_dim:] - - k1, k2 = k_rot.chunk(2, dim=-1) - c, s = rearrange(cos[:seqlen], "s d -> s 1 d"), rearrange( - sin[:seqlen], "s d -> s 1 d" - ) - k1, k2, c, s = [t.to(dtype=torch.float32) for t in [k1, k2, c, s]] - - k_rot = torch.cat([k1 * c - k2 * s, k1 * s + k2 * c], axis=-1).to(kv.dtype) - - return torch.cat( - [ - torch.cat([k_rot, k_pass], axis=-1).unsqueeze(2), - kv[:, :, 1:2, :, :], - ], - axis=2, - ) - - -# @torch.compile -def _apply_rotary_emb_qkv( - qkv: torch.FloatTensor, - cos: torch.FloatTensor, - sin: torch.FloatTensor, - cos_k: Optional[torch.FloatTensor] = None, - sin_k: Optional[torch.FloatTensor] = None, -) -> torch.FloatTensor: - _, seqlen, _, _, _ = qkv.shape - _, rotary_dim = cos.shape - rotary_dim *= 2 - - q_rot = qkv[:, :, 0, :, :rotary_dim] - q_pass = qkv[:, :, 0, :, rotary_dim:] - - k_rot = qkv[:, :, 1, :, :rotary_dim] - k_pass = qkv[:, :, 1, :, rotary_dim:] - - q1, q2 = q_rot.chunk(2, dim=-1) - k1, k2 = k_rot.chunk(2, dim=-1) - c, s = rearrange(cos[:seqlen], "s d -> s 1 d"), rearrange( - sin[:seqlen], "s d -> s 1 d" - ) - q1, q2, k1, k2, c, s = [t.to(dtype=torch.float32) for t in [q1, q2, k1, k2, c, s]] - - q_rot = torch.cat([q1 * c - q2 * s, q1 * s + q2 * c], axis=-1).to(qkv.dtype) - k_rot = torch.cat([k1 * c - k2 * s, k1 * s + k2 * c], axis=-1).to(qkv.dtype) - - return torch.cat( - [ - torch.cat([q_rot, q_pass], axis=-1).unsqueeze(2), - torch.cat([k_rot, k_pass], axis=-1).unsqueeze(2), - qkv[:, :, 2:3, :, :], - ], - axis=2, - ) - - -class RotaryEmbedding(nn.Module): - """Rotary positional embedding (RoPE). - - Reference: - RoFormer: Enhanced Transformer with Rotary Position Embedding. - https://arxiv.org/pdf/2104.09864.pdf. - - """ - - def __init__( - self, - dim: int, - base: int = 10000, - scale_base: Optional[float] = None, - pos_idx_in_fp32: bool = True, - max_position_embeddings: int = 2048, - device: Optional[str] = None, - **kwargs, - ) -> None: - super().__init__() - - if scale_base is not None: - raise NotImplementedError - - self.dim = dim - self.base = float(base) - self.scale_base = scale_base - self.pos_idx_in_fp32 = pos_idx_in_fp32 - self.max_position_embeddings = max_position_embeddings - self.device = device - - # Generate and save the inverse frequency buffer (non-trainable) - inv_freq = self._compute_inv_freq(device) - self.register_buffer("inv_freq", inv_freq, persistent=False) - - # Generate and save the scale buffer (non-trainable) - scale = ( - (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) - / (1.4 * dim) - if scale_base is not None - else None - ) - self.register_buffer("scale", scale, persistent=False) - - # Initialize cached attributes since ONNX can't rely on dynamic initialization - self._update_cos_sin_cache( - max_position_embeddings, device=device, dtype=torch.float32 - ) - - def _compute_inv_freq(self, device: Optional[str] = None) -> torch.FloatTensor: - return 1.0 / ( - self.base - ** ( - torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) - / self.dim - ) - ) - - def _update_cos_sin_cache( - self, - seqlen: int, - device: Optional[str] = None, - dtype: Optional[torch.dtype] = None, - ) -> None: - self._seq_len_cached = seqlen - - # fp32 is preferred since the output of `torch.arange` can be quite large - # and bf16 would lose a lot of precision - if self.pos_idx_in_fp32: - t = torch.arange(seqlen, device=device, dtype=torch.float32) - if self.inv_freq.dtype != torch.float32: - inv_freq = self._compute_inv_freq(device=device) - else: - inv_freq = self.inv_freq - else: - t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype) - inv_freq = self.inv_freq - - # `torch.outer` is preferred since `torch.einsum` converts from fp32 to fp16 if used with AMP - freqs = torch.outer(t, inv_freq) - if self.scale is None: - self._cos_cached = torch.cos(freqs).to(dtype) - self._sin_cached = torch.sin(freqs).to(dtype) - else: - power = ( - torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - - seqlen // 2 - ) / self.scale_base - scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1") - - # Force the scale multiplication to happen in fp32 - self._cos_cached = (torch.cos(freqs) * scale).to(dtype) - self._sin_cached = (torch.sin(freqs) * scale).to(dtype) - self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype) - self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype) - - def forward( - self, - qkv: torch.Tensor, - kv: Optional[torch.Tensor] = None, - seqlen_offset: int = 0, - **kwargs, - ) -> Tuple[torch.Tensor, torch.Tensor]: - if ( - self._seq_len_cached < qkv.shape[1] + seqlen_offset - or self._cos_cached.device != qkv.device - or self._cos_cached.dtype != qkv.dtype - or (self.training and self._cos_cached.is_inference()) - ): - self._update_cos_sin_cache( - qkv.shape[1] + seqlen_offset, device=qkv.device, dtype=qkv.dtype - ) - - if kv is None: - return _apply_rotary_emb_qkv( - qkv, - self._cos_cached[seqlen_offset:], - self._sin_cached[seqlen_offset:], - ) - else: - q = _apply_rotary_emb( - qkv, - self._cos_cached[seqlen_offset:], - self._sin_cached[seqlen_offset:], - ) - kv = _apply_rotary_emb_kv( - kv, - self._cos_cached[seqlen_offset:], - self._sin_cached[seqlen_offset:], - ) - - return q, kv - - -class MLP(nn.Module): - """Multi-Layer Perceptron. - - Reference: - Attention Is All You Need. - https://arxiv.org/pdf/1706.03762.pdf. - - """ - - def __init__( - self, - config: PretrainedConfig, - n_inner: Optional[int] = None, - act_fn: Optional[str] = None, - ) -> None: - super().__init__() - - act_fn = config.activation_function if act_fn is None else act_fn - - n_inner = getattr(config, "n_inner", None) if n_inner is None else n_inner - n_inner = n_inner if n_inner is not None else 4 * config.n_embd - - self.fc1 = nn.Linear(config.n_embd, n_inner) - self.fc2 = nn.Linear(n_inner, config.n_embd) - self.act = ACT2FN[act_fn] - - def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor: - hidden_states = self.fc1(hidden_states) - hidden_states = self.act(hidden_states) - hidden_states = self.fc2(hidden_states) - - return hidden_states - - -class SelfAttention(nn.Module): - """Self-attention layer (compatible with PyTorch). - - Reference: - https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py. - - """ - - def __init__( - self, - causal: bool = True, - softmax_scale: Optional[float] = None, - attention_dropout: float = 0.0, - ) -> None: - super().__init__() - - self.causal = causal - self.softmax_scale = softmax_scale - self.drop = nn.Dropout(attention_dropout) - - @torch.autocast("cpu", enabled=False) - @torch.autocast("cuda", enabled=False) - def forward( - self, - qkv: torch.FloatTensor, - causal: bool = None, - key_padding_mask: Optional[torch.BoolTensor] = None, - **kwargs, - ) -> torch.FloatTensor: - batch_size, seqlen = qkv.shape[0], qkv.shape[1] - q, k, v = qkv.unbind(dim=2) - - q = q.to(torch.float32) - k = k.to(torch.float32) - - causal = self.causal if causal is None else causal - softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1]) - - # Autocast is manually disabled to avoid `torch.einsum` performing the operation - # using float16, which might lead to overflow - scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale) - - if key_padding_mask is not None: - padding_mask = torch.full( - (batch_size, seqlen), -10000.0, dtype=scores.dtype, device=scores.device - ) - padding_mask.masked_fill_(key_padding_mask, 0.0) - - scores = scores + rearrange(padding_mask, "b s -> b 1 1 s") - - if causal: - causal_mask = torch.triu( - torch.full((seqlen, seqlen), -10000.0, device=scores.device), 1 - ) - scores = scores + causal_mask.to(dtype=scores.dtype) - - attention = torch.softmax(scores, dim=-1).to(v.dtype) - attention = self.drop(attention) - - output = torch.einsum("bhts,bshd->bthd", attention, v) - - return output - - -class CrossAttention(nn.Module): - """Cross-attention layer (compatible with PyTorch). - - Reference: - https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py. - - """ - - def __init__( - self, - causal: bool = True, - softmax_scale: Optional[float] = None, - attention_dropout: float = 0.0, - ) -> None: - super().__init__() - - self.causal = causal - self.softmax_scale = softmax_scale - self.drop = nn.Dropout(attention_dropout) - - @torch.autocast("cpu", enabled=False) - @torch.autocast("cuda", enabled=False) - def forward( - self, - q: torch.FloatTensor, - kv: torch.FloatTensor, - causal: bool = None, - key_padding_mask: Optional[torch.BoolTensor] = None, - **kwargs, - ) -> torch.FloatTensor: - batch_size, seqlen_q = q.shape[0], q.shape[1] - seqlen_k = kv.shape[1] - - if kv.shape[3] != q.shape[2]: - kv = repeat(kv, "... hkv d -> ... (hkv g) d", g=q.shape[2] // kv.shape[3]) - k, v = kv.unbind(dim=2) - - q = q.to(torch.float32) - k = k.to(torch.float32) - - causal = self.causal if causal is None else causal - softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1]) - - # Autocast is manually disabled to avoid `torch.einsum` performing the operation - # using float16, which might lead to overflow - scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale) - - if key_padding_mask is not None: - padding_mask = torch.full( - (batch_size, seqlen_k), - -10000.0, - dtype=scores.dtype, - device=scores.device, - ) - padding_mask.masked_fill_(key_padding_mask, 0.0) - - scores = scores + rearrange(padding_mask, "b s -> b 1 1 s") - - if causal: - rows = rearrange( - torch.arange(seqlen_q, device=q.device, dtype=torch.long), "s -> s 1" - ) - cols = torch.arange(seqlen_k, device=k.device, dtype=torch.long) - causal_mask = cols > rows + seqlen_k - seqlen_q - - scores = scores.masked_fill(causal_mask, -10000.0) - - attention = torch.softmax(scores, dim=-1).to(v.dtype) - attention = self.drop(attention) - - output = torch.einsum("bhts,bshd->bthd", attention, v) - - return output - - -def _find_mha_dims( - config: PretrainedConfig, - n_head: Optional[int] = None, - n_head_kv: Optional[int] = None, - head_dim: Optional[int] = None, -) -> Tuple[int, int]: - if n_head is None and head_dim is None: - head_dim = config.n_embd // config.n_head - n_head = config.n_head - elif n_head is None or head_dim is None: - raise ValueError("`n_head` and `head_dim` must be both specified or `None`.") - - if n_head_kv is None: - n_head_kv = getattr(config, "n_head_kv", None) or n_head - - return n_head, n_head_kv, head_dim - - -def _update_kv_cache( - kv: torch.FloatTensor, inference_params: InferenceParams, layer_idx: int -) -> torch.FloatTensor: - num_heads, head_dim = kv.shape[-2:] - - if layer_idx not in inference_params.key_value_memory_dict: - inference_params.key_value_memory_dict[layer_idx] = torch.empty( - inference_params.max_batch_size, - inference_params.max_seqlen, - 2, - num_heads, - head_dim, - dtype=kv.dtype, - device=kv.device, - ) - - batch_start = inference_params.batch_size_offset - batch_end = batch_start + kv.shape[0] - - sequence_start = inference_params.seqlen_offset - sequence_end = sequence_start + kv.shape[1] - - # When the current sequence length is equal to or larger than the maximum sequence length, - # we need to concatenate the current `kv` with the cached `kv` to expand its length - if sequence_end >= inference_params.max_seqlen: - inference_params.key_value_memory_dict[layer_idx] = torch.concatenate( - (inference_params.key_value_memory_dict[layer_idx], kv), dim=1 - ) - - inference_params.key_value_memory_dict[layer_idx][ - batch_start:batch_end, sequence_start:sequence_end, ... - ] = kv - kv = inference_params.key_value_memory_dict[layer_idx][ - batch_start:batch_end, :sequence_end, ... - ] - - return kv - - -class MHA(nn.Module): - """Multi-head attention layer.""" - - def __init__( - self, - config: PretrainedConfig, - dtype: Optional[torch.dtype] = None, - device: Optional[str] = None, - rotary_dim: Optional[int] = None, - rotary_base: float = 10000.0, - rotary_scale_base: Optional[float] = None, - n_head: Optional[int] = None, - n_head_kv: Optional[int] = None, - head_dim: Optional[int] = None, - bias: bool = True, - causal: bool = True, - softmax_scale: Optional[float] = None, - layer_idx: Optional[int] = None, - return_residual: bool = False, - checkpointing: bool = False, - ) -> None: - super().__init__() - - # Rotary embedding - self.rotary_dim = ( - rotary_dim if rotary_dim is not None else getattr(config, "rotary_dim", 0) - ) - - if self.rotary_dim > 0: - self.rotary_emb = RotaryEmbedding( - self.rotary_dim, - base=rotary_base, - scale_base=rotary_scale_base, - device=device, - max_position_embeddings=config.n_positions, - ) - - # MLP - self.n_head, self.n_head_kv, self.head_dim = _find_mha_dims( - config, n_head=n_head, n_head_kv=n_head_kv, head_dim=head_dim - ) - op_size = self.head_dim * (self.n_head + 2 * self.n_head_kv) - hidden_size = config.n_embd - - linear_cls = FusedDense if config.fused_dense else nn.Linear - if linear_cls is None: - linear_cls = nn.Linear - - self.Wqkv = linear_cls( - hidden_size, op_size, bias=bias, device=device, dtype=dtype - ) - self.out_proj = linear_cls( - hidden_size, hidden_size, bias=bias, device=device, dtype=dtype - ) - - # Attention - self.inner_attn = SelfAttention( - causal=causal, - softmax_scale=softmax_scale, - attention_dropout=config.attn_pdrop, - ) - self.inner_cross_attn = CrossAttention( - causal=causal, - softmax_scale=softmax_scale, - attention_dropout=config.attn_pdrop, - ) - - self.layer_idx = layer_idx - self.return_residual = return_residual - self.checkpointing = checkpointing - - def _forward_self_attn( - self, x: torch.FloatTensor, key_padding_mask: Optional[torch.BoolTensor] - ) -> torch.FloatTensor: - qkv = self.Wqkv(x) - qkv = rearrange( - qkv, "... (three h d) -> ... three h d", three=3, d=self.head_dim - ) - - if self.rotary_dim > 0: - qkv = self.rotary_emb(qkv) - - if self.checkpointing: - return torch.utils.checkpoint.checkpoint( - self.inner_attn, qkv, key_padding_mask=key_padding_mask - ) - - return self.inner_attn(qkv, key_padding_mask=key_padding_mask) - - def _forward_cross_attn( - self, - x: torch.FloatTensor, - past_key_values: Optional[InferenceParams], - key_padding_mask: Optional[torch.BoolTensor], - ) -> torch.FloatTensor: - batch_size = x.shape[0] - - qkv = self.Wqkv(x) - - q = qkv[..., : self.n_head * self.head_dim] - q = rearrange(q, "... (h d) -> ... h d", d=self.head_dim) - - kv = qkv[..., self.n_head * self.head_dim :] - kv = rearrange(kv, "... (two hkv d) -> ... two hkv d", two=2, d=self.head_dim) - - seqlen_offset = ( - past_key_values.seqlen_offset if past_key_values is not None else 0 - ) - causal = None if seqlen_offset == 0 else False - if self.rotary_dim > 0: - q, kv = self.rotary_emb(q, kv=kv, seqlen_offset=seqlen_offset) - - if past_key_values is not None: - kv = _update_kv_cache(kv, past_key_values, self.layer_idx) - - if self.checkpointing: - return torch.utils.checkpoint.checkpoint( - self.inner_cross_attn, - q, - kv, - key_padding_mask=key_padding_mask, - causal=causal, - ) - - return self.inner_cross_attn( - q, kv, key_padding_mask=key_padding_mask, causal=causal - ) - - def forward( - self, - x: torch.FloatTensor, - past_key_values: Optional[InferenceParams] = None, - attention_mask: Optional[Union[torch.LongTensor, torch.BoolTensor]] = None, - **kwargs, - ) -> Tuple[torch.FloatTensor, torch.FloatTensor]: - if attention_mask is not None: - attention_mask = attention_mask.bool() - else: - attention_mask = None - - # MHA - if self.n_head == self.n_head_kv: - if past_key_values is None: - # If `past_key_values` are not supplied, we run self-attention - attn_output = self._forward_self_attn(x, attention_mask) - else: - # If `past_key_values` are supplied, it means that we might have cached values and - # could take advantage of cross-attention - attn_output = self._forward_cross_attn( - x, past_key_values, attention_mask - ) - # MQA / GQA - else: - # Regardless of `past_key_values` being supplied or not, it always use cross-attention - # because `q` and `kv` lengths might be different - attn_output = self._forward_cross_attn(x, past_key_values, attention_mask) - - output = rearrange(attn_output, "... h d -> ... (h d)") - output = self.out_proj(output) - - return output if not self.return_residual else (output, x) - - -class ParallelBlock(nn.Module): - """Parallel block. - - This block applies parallel mixer and MLP layers to the input (used in GPT-J and CodeGen). - - """ - - def __init__( - self, - config: PretrainedConfig, - block_idx: Optional[int] = None, - ) -> None: - super().__init__() - - self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon) - self.resid_dropout = nn.Dropout(config.resid_pdrop) - self.block_idx = block_idx - - self.mixer = MHA(config, layer_idx=block_idx) - self.mlp = MLP(config) - - def forward( - self, - hidden_states: torch.FloatTensor, - past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None, - attention_mask: Optional[torch.BoolTensor] = None, - **kwargs, - ) -> torch.FloatTensor: - residual = hidden_states - hidden_states = self.ln(hidden_states) - - attn_outputs = self.mixer( - hidden_states, - past_key_values=past_key_values, - attention_mask=attention_mask, - ) - if isinstance(attn_outputs, tuple): - attn_outputs = attn_outputs[0] - - attn_outputs = self.resid_dropout(attn_outputs) - feed_forward_hidden_states = self.resid_dropout(self.mlp(hidden_states)) - - hidden_states = attn_outputs + feed_forward_hidden_states + residual - - return hidden_states - - -class CausalLMHead(nn.Module): - """Causal Language Modeling head. - - Reference: - Improving Language Understanding by Generative Pre-Training. - https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf. - - """ - - def __init__(self, config: PretrainedConfig) -> None: - super().__init__() - - self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon) - self.linear = nn.Linear(config.n_embd, config.vocab_size) - - def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor: - hidden_states = self.ln(hidden_states) - logits = self.linear(hidden_states).to(torch.float32) - - return logits - - -class CausalLMLoss(nn.Module): - """Causal Language Modeling loss. - - Reference: - Improving Language Understanding by Generative Pre-Training. - https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf. - - """ - - def __init__(self, shift_labels: bool = True) -> None: - super().__init__() - - self.shift_labels = shift_labels - self.loss_fct = nn.CrossEntropyLoss() - - def forward( - self, logits: torch.FloatTensor, labels: torch.LongTensor - ) -> torch.FloatTensor: - if self.shift_labels: - logits = logits[..., :-1, :].contiguous() - labels = labels[..., 1:].contiguous() - - loss = self.loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1)) - - return loss - - -class PhiPreTrainedModel(PreTrainedModel): - """Phi pre-trained model.""" - - config_class = PhiConfig - base_model_prefix = "transformer" - supports_gradient_checkpointing = False - _no_split_modules = ["ParallelBlock"] - - def __init__(self, *inputs, **kwargs) -> None: - super().__init__(*inputs, **kwargs) - - def prepare_inputs_for_generation( - self, - input_ids: torch.LongTensor = None, - inputs_embeds: torch.FloatTensor = None, - past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None, - attention_mask: Optional[Union[torch.LongTensor, torch.BoolTensor]] = None, - **kwargs, - ) -> Dict[str, Any]: - if inputs_embeds is not None: - max_batch_size = inputs_embeds.shape[0] - seqlen_offset = inputs_embeds.shape[1] + input_ids.shape[1] - 2 - elif input_ids is not None: - max_batch_size = input_ids.shape[0] - seqlen_offset = input_ids.shape[1] - 1 - else: - raise ValueError( - "You have to specify either `input_ids` or `inputs_embeds`." - ) - - args = {} - - if past_key_values is None or not ( - isinstance(past_key_values, InferenceParams) - ): - past_key_values = InferenceParams( - max_seqlen=self.config.n_positions, - max_batch_size=max_batch_size, - seqlen_offset=0, - batch_size_offset=0, - key_value_memory_dict={}, - lengths_per_sample=None, - ) - if inputs_embeds is not None: - args = {"inputs_embeds": inputs_embeds} - elif input_ids is not None: - args = {"input_ids": input_ids} - else: - raise ValueError( - "You have to specify either `input_ids` or `inputs_embeds`." - ) - else: - # Assume that `past_key_values` has cached all tokens up to the last token in `input_ids` - past_key_values.seqlen_offset = seqlen_offset - input_ids = input_ids[:, -1].unsqueeze(-1) - args = {"input_ids": input_ids} - - return { - **args, - "past_key_values": past_key_values, - "attention_mask": attention_mask, - } - - -class PhiModel(PhiPreTrainedModel): - """Phi model.""" - - _keys_to_ignore_on_load_missing = [""] - _keys_to_ignore_on_load_unexpected = [r"h\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)"] - - def __init__(self, config: PhiConfig) -> None: - super().__init__(config) - - self.embd = Embedding(config) - self.h = nn.ModuleList( - [ParallelBlock(config, block_idx=i) for i in range(config.n_layer)] - ) - self.gradient_checkpointing = config.gradient_checkpointing - self.post_init() - - def get_input_embeddings(self) -> nn.Embedding: - return self.embd.wte - - def set_input_embeddings(self, new_embeddings: nn.Embedding) -> None: - self.embd.wte = new_embeddings - - def forward( - self, - input_ids: torch.LongTensor = None, - inputs_embeds: torch.FloatTensor = None, - past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None, - attention_mask: Optional[torch.BoolTensor] = None, - ) -> torch.FloatTensor: - if input_ids is not None and inputs_embeds is not None: - raise ValueError( - "You cannot specify both `input_ids` and `inputs_embeds` at the same time." - ) - elif input_ids is None and inputs_embeds is None: - raise ValueError( - "You have to specify either `input_ids` or `inputs_embeds`." - ) - elif input_ids is not None: - hidden_states = self.embd(input_ids) - else: - hidden_states = inputs_embeds - - for layer in self.h: - if self.gradient_checkpointing: - hidden_states = torch.utils.checkpoint.checkpoint( - layer.__call__, - hidden_states, - past_key_values, - attention_mask, - use_reentrant=True, - ) - else: - hidden_states = layer( - hidden_states, - past_key_values=past_key_values, - attention_mask=attention_mask, - ) - - return hidden_states - - -class PhiForCausalLM(PhiPreTrainedModel): - """Phi for Causal Language Modeling.""" - - _keys_to_ignore_on_load_missing = [""] - _keys_to_ignore_on_load_unexpected = [ - r"transformer\.h\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)" - ] - - def __init__(self, config: PhiConfig) -> None: - super().__init__(config) - - self.transformer = PhiModel(config) - self.lm_head = CausalLMHead(config) - self.loss = CausalLMLoss() - - self.post_init() - - def get_output_embeddings(self) -> nn.Linear: - return self.lm_head.linear - - def set_output_embeddings(self, new_embeddings: nn.Linear) -> None: - self.lm_head.linear = new_embeddings - - def forward( - self, - input_ids: torch.LongTensor = None, - inputs_embeds: torch.FloatTensor = None, - past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None, - attention_mask: Optional[torch.BoolTensor] = None, - labels: Optional[torch.LongTensor] = None, - **kwargs, - ) -> CausalLMOutputWithPast: - hidden_states = self.transformer( - input_ids, - inputs_embeds, - past_key_values=past_key_values, - attention_mask=attention_mask, - ) - lm_logits = self.lm_head(hidden_states) - - loss = None - if labels is not None: - loss = self.loss(lm_logits, labels) - - return CausalLMOutputWithPast( - loss=loss, logits=lm_logits, past_key_values=past_key_values - ) diff --git a/moondream/text_model.py b/moondream/text_model.py index 7ac99a0..184eae8 100644 --- a/moondream/text_model.py +++ b/moondream/text_model.py @@ -1,91 +1,19 @@ -import torch +from torch import nn import transformers -from transformers import CodeGenTokenizerFast as Tokenizer -from accelerate import init_empty_weights, load_checkpoint_and_dispatch -from .phi.configuration_phi import PhiConfig -from .phi.modeling_phi import PhiForCausalLM -import re +from .modeling_phi import PhiForCausalLM +from .configuration_moondream import PhiConfig transformers.logging.set_verbosity_error() -class TextModel: - def __init__(self, model_path: str = "model") -> None: +class TextModel(nn.Module): + def __init__(self, config) -> None: super().__init__() - # Determine if CUDA (GPU) is available and use it; otherwise, use CPU - self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - self.tokenizer = Tokenizer.from_pretrained(f"{model_path}/tokenizer") - phi_config = PhiConfig.from_pretrained(f"{model_path}/text_model_cfg.json") - - with init_empty_weights(): - self.model = PhiForCausalLM(phi_config) - - self.model = load_checkpoint_and_dispatch( - self.model, - f"{model_path}/text_model.pt", - device_map={"": self.device.type}, - ).half() - - self.text_emb = self.model.get_input_embeddings().to(self.device) - - def input_embeds(self, prompt, image_embeds): - embeds = [] - - def _add_toks(toks): - embeds.append(self.text_emb(toks)) - - def _tokenize(txt): - return self.tokenizer( - txt, return_tensors="pt", add_special_tokens=False - ).input_ids.to(self.model.device) - - # Add BOS token - _add_toks( - torch.tensor([[self.tokenizer.bos_token_id]], device=self.model.device) - ) - - if "" not in prompt: - embeds.append(self.text_emb(_tokenize(prompt))) + if type(config.phi_config) == dict: + phi_config = PhiConfig(**config.phi_config) else: - assert prompt.count("") == 1 - before, after = prompt.split("") - embeds.append(self.text_emb(_tokenize(f"{before}"))) - embeds.append(image_embeds.to(self.model.device)) - embeds.append(self.text_emb(_tokenize(f"{after}"))) + phi_config = config.phi_config - return torch.cat(embeds, dim=1) - - def generate( - self, image_embeds, prompt, eos_text="Human:", max_new_tokens=128, **kwargs - ): - eos_tokens = self.tokenizer(eos_text, add_special_tokens=False)[0].ids - - generate_config = { - "eos_token_id": eos_tokens, - "bos_token_id": self.tokenizer.bos_token_id, - "pad_token_id": self.tokenizer.eos_token_id, - "max_new_tokens": max_new_tokens, - **kwargs, - } - - with torch.no_grad(): - inputs_embeds = self.input_embeds(prompt, image_embeds) - output_ids = self.model.generate( - inputs_embeds=inputs_embeds, **generate_config - ) - - return self.tokenizer.batch_decode(output_ids, skip_special_tokens=True) - - def answer_question(self, image_embeds, question, **kwargs): - prompt = f"\n\nQuestion: {question}\n\nAnswer:" - answer = self.generate( - image_embeds, - prompt, - eos_text="", - max_new_tokens=128, - **kwargs, - )[0] - - return re.sub("<$", "", re.sub("END$", "", answer)).strip() + self.model = PhiForCausalLM(phi_config) + self.text_emb = self.model.get_input_embeddings() diff --git a/moondream/util.py b/moondream/util.py new file mode 100644 index 0000000..d4b2a31 --- /dev/null +++ b/moondream/util.py @@ -0,0 +1,13 @@ +import torch + + +def detect_device(): + """ + Detects the appropriate device to run on, and return the device and dtype. + """ + if torch.cuda.is_available(): + return torch.device("cuda"), torch.float16 + elif torch.backends.mps.is_available(): + return torch.device("mps"), torch.float16 + else: + return torch.device("cpu"), torch.float32 diff --git a/moondream/vision_encoder.py b/moondream/vision_encoder.py index df4f6ab..e22fbca 100644 --- a/moondream/vision_encoder.py +++ b/moondream/vision_encoder.py @@ -1,4 +1,5 @@ import torch +from torch import nn from PIL import Image from einops import rearrange from torchvision.transforms.v2 import ( @@ -9,28 +10,141 @@ from torchvision.transforms.v2 import ( ToDtype, Normalize, ) +import timm -class VisionEncoder: - def __init__(self, model_path: str = "model") -> None: - # Determine if CUDA (GPU) is available and use it; otherwise, use CPU - self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - self.model = torch.jit.load(f"{model_path}/vision.pt").to(self.device).to(dtype=torch.float16) +class VisualHolder(nn.Module): + def __init__(self, model): + super().__init__() + self.visual = model + + def forward(self, x): + return self.visual(x) + + +class ModelHolder(nn.Module): + def __init__(self, model): + super().__init__() + self.model = model + + def forward(self, x): + return self.model(x) + + +class LinearPatchEmbedding(nn.Module): + def __init__(self, conv): + super().__init__() + self.linear = nn.Linear(588, 1152) + self.linear.weight.data = conv.weight.data.view(1152, -1) + if conv.bias is not None: + self.linear.bias.data = conv.bias.data + + def forward(self, x): + return self.linear(x) + + +class MLP(nn.Module): + def __init__( + self, + in_features: int, + hidden_features: int = None, + out_features: int = None, + act_layer: nn.Module = nn.GELU, + ) -> None: + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + + torch.nn.init.kaiming_normal_( + self.fc1.weight, mode="fan_in", nonlinearity="relu" + ) + torch.nn.init.kaiming_normal_( + self.fc2.weight, mode="fan_in", nonlinearity="relu" + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.fc1(x) + x = self.act(x) + x = self.fc2(x) + return x + + +class VisionProjection(nn.Module): + def __init__(self): + super().__init__() + + image_embedding_dim = 1152 + model_dim = 2048 + hidden_dim = model_dim * 4 + + self.mlp1 = MLP(image_embedding_dim, hidden_dim, model_dim) + self.mlp2 = MLP(model_dim, hidden_dim, model_dim) + self.ln = nn.LayerNorm(model_dim) + + @property + def device(self): + return self.mlp1.fc1.weight.device + + def forward(self, x): + x = self.mlp1(x) + x = self.ln(x) + x = x + self.mlp2(x) + return x + + +class VisionTower(nn.Module): + def __init__(self): + super().__init__() + + self.encoder = ModelHolder( + VisualHolder(timm.create_model("vit_so400m_patch14_siglip_384")) + ) + self.encoder.model.visual.patch_embed = LinearPatchEmbedding( + self.encoder.model.visual.patch_embed.proj + ) + self.encoder.model.visual.attn_pool = nn.Identity() + + self.projection = VisionProjection() + + def forward(self, x): + x = self.encoder(x) + x = self.projection(x) + return x + + +class VisionEncoder(nn.Module): + def __init__(self) -> None: + super().__init__() + + self.model = VisionTower() self.preprocess = Compose( [ - Resize(size=(384, 384), interpolation=InterpolationMode.BICUBIC), + Resize(size=(378, 378), interpolation=InterpolationMode.BICUBIC), ToImage(), - ToDtype(torch.float16, scale=True), + ToDtype(torch.float32, scale=True), Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), ] ) + @property + def device(self): + return self.model.projection.mlp1.fc1.weight.device + + @property + def dtype(self): + return self.model.projection.mlp1.fc1.weight.dtype + def __call__(self, image: Image) -> torch.Tensor: with torch.no_grad(): - image_vec = self.preprocess(image.convert("RGB")).unsqueeze(0).to(self.device) - image_vec = image_vec[:, :, :-6, :-6] + image_vec = ( + self.preprocess(image.convert("RGB")) + .unsqueeze(0) + .to(self.device, dtype=self.dtype) + ) image_vec = rearrange( image_vec, "b c (h p1) (w p2) -> b (h w) (c p1 p2)", p1=14, p2=14 ) - return self.model(image_vec) diff --git a/sample.py b/sample.py index 2333c2f..711db8f 100644 --- a/sample.py +++ b/sample.py @@ -1,43 +1,77 @@ -from moondream import VisionEncoder, TextModel -from PIL import Image -from huggingface_hub import snapshot_download +import torch import argparse +from PIL import Image +from moondream import Moondream, detect_device +from queue import Queue from threading import Thread -from transformers import TextIteratorStreamer +from transformers import ( + TextIteratorStreamer, + CodeGenTokenizerFast as Tokenizer, +) import re -model_path = snapshot_download("vikhyatk/moondream1") +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--image", type=str, required=True) + parser.add_argument("--prompt", type=str, required=False) + parser.add_argument("--cpu", action="store_true") + args = parser.parse_args() -vision_encoder = VisionEncoder(model_path) -text_model = TextModel(model_path) + if args.cpu: + device = torch.device("cpu") + dtype = torch.float32 + else: + device, dtype = detect_device() + if device != torch.device("cpu"): + print("Using device:", device) + print("If you run into issues, pass the `--cpu` flag to this script.") + print() -parser = argparse.ArgumentParser() -parser.add_argument("--image", type=str, required=True) -parser.add_argument("--prompt", type=str, required=False) -args = parser.parse_args() + image_path = args.image + prompt = args.prompt -image = Image.open(args.image) -image_embeds = vision_encoder(image) + model_id = "vikhyatk/moondream1" + tokenizer = Tokenizer.from_pretrained(model_id) + moondream = Moondream.from_pretrained(model_id).to(device=device, dtype=dtype) + moondream.eval() -if args.prompt is None: - while True: - question = input("> ") + image = Image.open(image_path) + image_embeds = moondream.encode_image(image) - streamer = TextIteratorStreamer(text_model.tokenizer, skip_special_tokens=True) - generation_kwargs = dict( - image_embeds=image_embeds, question=question, streamer=streamer - ) - thread = Thread(target=text_model.answer_question, kwargs=generation_kwargs) - thread.start() + if prompt is None: + chat_history = "" - buffer = "" - for new_text in streamer: - buffer += new_text - if not new_text.endswith("<") and not new_text.endswith("END"): - print(buffer, end="", flush=True) - buffer = "" - print(re.sub("<$", "", re.sub("END$", "", buffer))) -else: - question = args.prompt - print(">", question) - print(text_model.answer_question(image_embeds, question)) + while True: + question = input("> ") + + result_queue = Queue() + + streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True) + + # Separate direct arguments from keyword arguments + thread_args = (image_embeds, question, tokenizer, chat_history) + thread_kwargs = {"streamer": streamer, "result_queue": result_queue} + + thread = Thread( + target=moondream.answer_question, + args=thread_args, + kwargs=thread_kwargs, + ) + thread.start() + + buffer = "" + for new_text in streamer: + buffer += new_text + if not new_text.endswith("<") and not new_text.endswith("END"): + print(buffer, end="", flush=True) + buffer = "" + print(re.sub("<$", "", re.sub("END$", "", buffer))) + + thread.join() + + answer = result_queue.get() + chat_history += f"Question: {question}\n\nAnswer: {answer}\n\n" + else: + print(">", prompt) + answer = moondream.answer_question(image_embeds, prompt, tokenizer) + print(answer)