github
/
minimind-o
mirror of https://github.com/jingyaogong/minimind-o.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge 03b593621b into f900448c60

pull/21/merge
Sihua ZHANG 2026-09-23 18:43:35 +08:00 committed by GitHub
parent f900448c60 03b593621b
commit 4c9755ea49
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 13 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
model/model_omni.py
View File

@ -1,5 +1,7 @@
import os, math, torch, soundfile as sf, librosa, warnings, numpy as np, onnxruntime as ort, logging, contextlib, io import os, math, torch, soundfile as sf, librosa, warnings, numpy as np, onnxruntime as ort, logging, contextlib, io
from types import SimpleNamespace from types import SimpleNamespace
from dataclasses import dataclass
from typing import Optional
from torch import nn from torch import nn
from torch.nn import functional as F from torch.nn import functional as F
from transformers.modeling_outputs import MoeCausalLMOutputWithPast from transformers.modeling_outputs import MoeCausalLMOutputWithPast
@ -28,6 +30,11 @@ class OmniConfig(MiniMindConfig):
self.image_token_len = kwargs.get("image_token_len", 64) self.image_token_len = kwargs.get("image_token_len", 64)
self.bridge_layer = kwargs.get("bridge_layer", self.num_hidden_layers // 2 - 1) self.bridge_layer = kwargs.get("bridge_layer", self.num_hidden_layers // 2 - 1)
# 避免 DDP + compile 重建 OmniCausalLMOutputWithPast（不支持 audio_logits 参数）
@dataclass
class OmniCausalLMOutputWithPast(MoeCausalLMOutputWithPast):
audio_logits: Optional[list[torch.FloatTensor]] = None
class MMAudioProjector(nn.Module): class MMAudioProjector(nn.Module):
def __init__(self, in_dim, out_dim): def __init__(self, in_dim, out_dim):
super().__init__() super().__init__()
@ -311,9 +318,12 @@ class MiniMindOmni(MiniMindForCausalLM):
text_logits = self.thinker.lm_head(h_thinker[:, slice_indices, :]) text_logits = self.thinker.lm_head(h_thinker[:, slice_indices, :])
audio_logits = self.talker.lm_head(h_talker[:, slice_indices, :]) audio_logits = self.talker.lm_head(h_talker[:, slice_indices, :])
out = MoeCausalLMOutputWithPast(aux_loss=aux_loss, logits=text_logits, past_key_values=presents) return OmniCausalLMOutputWithPast(
out.audio_logits = audio_logits aux_loss=aux_loss,
return out logits=text_logits,
past_key_values=presents,
audio_logits=audio_logits,
)
@torch.inference_mode() @torch.inference_mode()
def generate(self, input_ids, eos_token_id=2, max_new_tokens=1024, temperature=0.75, top_p=0.90, def generate(self, input_ids, eos_token_id=2, max_new_tokens=1024, temperature=0.75, top_p=0.90,