Upload 3 files

Qwenov3Config.py

@@ -0,0 +1,186 @@
+from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin
+from modelscope import AutoConfig, AutoProcessor, AutoModel, AutoTokenizer, AutoModelForCausalLM
+import torch
+import torch.nn as nn
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+class Qwenov3Config(PretrainedConfig):
+    model_type = "Qwenov3"
+
+    def __init__(self, llm_model_path='Qwen/Qwen3-0.6B',
+                 vision_model_path='facebook/dinov3-vitl16-pretrain-lvd1689m',
+                 freeze_vision_model=False,
+                 freeze_llm_model=False,
+                 image_pad_num=49,
+                 training_scratch=False,
+                 num_hidden_layers=None,
+                 hidden_size=None,
+                 num_attention_heads=None,
+                 vocab_size=None,
+                 **kwargs):
+        self.vision_model_path = vision_model_path
+        self.llm_model_path = llm_model_path
+        self.freeze_vision_model = freeze_vision_model
+        self.freeze_llm_model = freeze_llm_model
+        self.image_pad_num = image_pad_num
+        self.freeze_vision_model = freeze_vision_model
+        self.training_scratch = training_scratch
+        self.num_hidden_layers = num_hidden_layers
+        self.hidden_size = hidden_size
+        self.num_attention_heads = num_attention_heads
+        self.vocab_size = vocab_size
+        
+        super().__init__(**kwargs)
+
+
+class Qwenov3(GenerationMixin, PreTrainedModel):
+    config_class = Qwenov3Config
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        if self.config.training_scratch:
+            self.vision_model = AutoModel.from_pretrained(self.config.vision_model_path, low_cpu_mem_usage=True,
+                                                          dtype=torch.bfloat16, attn_implementation="flash_attention_2")
+            self.llm_model = AutoModelForCausalLM.from_pretrained(self.config.llm_model_path, low_cpu_mem_usage=True,
+                                                                  dtype=torch.bfloat16,
+                                                                  attn_implementation="flash_attention_2")
+        else:
+            vision_config = AutoConfig.from_pretrained(self.config.vision_model_path)
+            self.vision_model = AutoModel.from_config(vision_config, attn_implementation="sdpa", dtype=torch.bfloat16)
+            llm_config = AutoConfig.from_pretrained(self.config.llm_model_path)
+            self.llm_model = AutoModelForCausalLM.from_config(llm_config, attn_implementation="sdpa", dtype=torch.bfloat16)
+
+        if self.config.num_hidden_layers is None:
+            self.config.num_hidden_layers = self.llm_model.config.num_hidden_layers
+        if self.config.hidden_size is None:
+            self.config.hidden_size = self.llm_model.config.hidden_size
+        if self.config.num_attention_heads is None:
+            self.config.num_attention_heads = self.llm_model.config.num_attention_heads
+        if self.config.vocab_size is None:
+            self.config.vocab_size = self.llm_model.config.vocab_size
+
+        self.processor = AutoProcessor.from_pretrained(self.config.vision_model_path)
+        self.tokenizer = AutoTokenizer.from_pretrained(self.config.llm_model_path, use_fast=True)
+
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+        if '<|image_pad|>' not in self.tokenizer.get_vocab():
+            self.tokenizer.add_tokens(['<|image_pad|>'])
+            self.llm_model.resize_token_embeddings(len(self.tokenizer), mean_resizing=True)
+        if '<|vision_start|>' not in self.tokenizer.get_vocab():
+            self.tokenizer.add_tokens(['<|vision_start|>'])
+            self.llm_model.resize_token_embeddings(len(self.tokenizer), mean_resizing=True)
+        if '<|vision_end|>' not in self.tokenizer.get_vocab():
+            self.tokenizer.add_tokens(['<|vision_end|>'])
+            self.llm_model.resize_token_embeddings(len(self.tokenizer), mean_resizing=True)
+
+        self.adapter = nn.Sequential(
+            nn.RMSNorm(4096, dtype=torch.bfloat16),
+            nn.Linear(4096, self.llm_model.config.hidden_size, dtype=torch.bfloat16),
+            nn.GELU(),
+            nn.Linear(self.llm_model.config.hidden_size, self.llm_model.config.hidden_size, dtype=torch.bfloat16)
+        )
+
+        if self.config.freeze_vision_model:
+            for param in self.vision_model.parameters():
+                param.requires_grad = False
+        if self.config.freeze_llm_model:
+            for param in self.llm_model.parameters():
+                param.requires_grad = False
+
+    def forward(self, input_ids=None, labels=None, pixel_values=None, attention_mask=None, 
+                inputs_embeds=None, past_key_values=None, use_cache=None, **kwargs):
+
+        if inputs_embeds is None:
+            text_embeds = self.llm_model.get_input_embeddings()(input_ids)
+            if pixel_values is not None:
+                image_embeds = self.vision_model(pixel_values).last_hidden_state
+                patch_embeds = image_embeds[:, 5:, :]  # [batch, 196, 1024]
+                b, num_patches, hidden_dim = patch_embeds.shape
+                patch_embeds = patch_embeds.view(b, num_patches // 4, hidden_dim * 4)  # [batch, 49, 4096]
+                image_features = self.adapter(patch_embeds)
+                text_embeds = text_embeds.to(image_features.dtype)
+                inputs_embeds = self.merge_input_ids_with_image_features(image_features, text_embeds, input_ids)
+            else:
+                inputs_embeds = text_embeds
+
+        outputs = self.llm_model(
+            inputs_embeds=inputs_embeds, 
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            return_dict=True
+        )
+        
+        logits = outputs.logits
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss(ignore_index=self.tokenizer.pad_token_id)
+            loss = loss_fct(
+                logits.view(-1, logits.size(-1)), labels.view(-1).to(logits.device)
+            )
+        
+        return CausalLMOutputWithPast(
+            loss=loss, 
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions
+        )
+
+    @torch.inference_mode()
+    def generate(self, input_ids=None, pixel_values=None, attention_mask=None, 
+                 max_new_tokens=512, temperature=0.7, top_p=0.8, top_k=20,
+                 do_sample=True, num_beams=1, use_cache=True, **kwargs):
+        if pixel_values is not None:
+            text_embeds = self.llm_model.get_input_embeddings()(input_ids)
+            image_embeds = self.vision_model(pixel_values).last_hidden_state
+            patch_embeds = image_embeds[:, 5:, :]
+            b, num_patches, hidden_dim = patch_embeds.shape
+            patch_embeds = patch_embeds.view(b, num_patches // 4, hidden_dim * 4)
+            image_features = self.adapter(patch_embeds)
+            text_embeds = text_embeds.to(image_features.dtype)
+            inputs_embeds = self.merge_input_ids_with_image_features(image_features, text_embeds, input_ids)
+            return self.llm_model.generate(
+                input_ids=input_ids,
+                inputs_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                top_k=top_k,
+                do_sample=do_sample,
+                num_beams=num_beams,
+                use_cache=use_cache,
+                pad_token_id=self.tokenizer.pad_token_id,
+                eos_token_id=self.tokenizer.eos_token_id,
+                **kwargs
+            )
+        else:
+            return self.llm_model.generate(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                top_k=top_k,
+                do_sample=do_sample,
+                num_beams=num_beams,
+                use_cache=use_cache,
+                pad_token_id=self.tokenizer.pad_token_id,
+                eos_token_id=self.tokenizer.eos_token_id,
+                **kwargs
+            )
+
+    def can_generate(self):
+        return True
+
+    def merge_input_ids_with_image_features(self, image_features, inputs_embeds, input_ids):
+        num_images, num_image_patches, embed_dim = image_features.shape
+        batch_indices, image_indices = torch.where(input_ids == self.tokenizer('<|image_pad|>')['input_ids'][0])
+        if len(batch_indices) == 0:
+            return inputs_embeds
+        inputs_embeds[batch_indices, image_indices] = image_features.view(-1, embed_dim)
+        return inputs_embeds

StreamlitUI.py

@@ -0,0 +1,204 @@
+import streamlit as st
+import torch
+from transformers import AutoModelForCausalLM, TextIteratorStreamer, AutoConfig
+import gc
+from threading import Thread
+from Qwenov3Config import Qwenov3Config, Qwenov3
+from PIL import Image
+
+MODEL_MAPPING = {
+    'QwenoV3-Pretrain': '',
+    'QwenoV3-SFT': '',
+}
+
+
+def unload_model():
+    if 'model' in st.session_state:
+        del st.session_state.model
+    if 'tokenizer' in st.session_state:
+        del st.session_state.tokenizer
+    if 'processor' in st.session_state:
+        del st.session_state.processor
+    if 'streamer' in st.session_state:
+        del st.session_state.streamer
+    torch.cuda.empty_cache()
+    gc.collect()
+
+
+def call_model(info_placeholder, messages, generated_text, message_placeholder, image=None):
+    info_placeholder.markdown(f'已选择{st.session_state.model_display}执行任务')
+    if image is not None:
+        image = Image.open(image).convert('RGB')
+        if '<image>' not in messages[1]['content']:
+            messages[1]['content'] = '<image>\n' + messages[1]['content']
+
+    query_text = st.session_state.tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+        enable_thinking=False
+    )
+    if '<image>' in query_text:
+        query_text = query_text.replace('<image>', '<|vision_start|>' + '<|image_pad|>' *
+                                        st.session_state.model.config.image_pad_num + '<|vision_end|>')
+    text_inputs = st.session_state.tokenizer(query_text, return_tensors="pt")
+    input_ids = text_inputs['input_ids'].to(st.session_state.model.device)
+    attention_mask = text_inputs['attention_mask'].to(st.session_state.model.device)
+    text_embeds = st.session_state.model.llm_model.get_input_embeddings()(input_ids)
+
+    if image is not None:
+        pixel_values = st.session_state.processor(images=image, return_tensors="pt")['pixel_values'].to(
+            st.session_state.model.device)
+        image_embeds = st.session_state.model.vision_model(pixel_values).last_hidden_state
+        patch_embeds = image_embeds[:, 5:, :]
+        b, num_patches, hidden_dim = patch_embeds.shape
+        patch_embeds = patch_embeds.view(b, num_patches // 4, hidden_dim * 4)
+        image_features = st.session_state.model.adapter(patch_embeds)
+        text_embeds = text_embeds.to(image_features.dtype)
+        inputs_embeds = st.session_state.model.merge_input_ids_with_image_features(image_features, text_embeds, input_ids)
+    else:
+        inputs_embeds = text_embeds
+
+    generate_params = dict(
+        inputs_embeds=inputs_embeds,
+        attention_mask=attention_mask,
+        max_new_tokens=st.session_state.max_new_tokens,
+        min_new_tokens=st.session_state.min_new_tokens,
+        do_sample=True,
+        temperature=st.session_state.temperature,
+        top_k=st.session_state.top_k,
+        top_p=st.session_state.top_p,
+        min_p=0.0,
+        repetition_penalty=st.session_state.repetition_penalty,
+        streamer=st.session_state.streamer,
+        eos_token_id=st.session_state.tokenizer.eos_token_id
+    )
+    thread = Thread(target=st.session_state.model.llm_model.generate, kwargs=generate_params)
+    thread.start()
+
+    for new_text in st.session_state.streamer:
+        generated_text += new_text
+        message_placeholder.markdown(generated_text)
+
+    return generated_text
+
+
+def ini_message():
+    if 'messages' not in st.session_state:
+        st.session_state.messages = [
+            {"role": "system", "content": "You are QwenoV3, a helpful assistant created by 天烨."},
+        ]
+    if 'uploaded_image' not in st.session_state:
+        st.session_state.uploaded_image = None
+
+
+def parameter_settings():
+    with st.sidebar:
+        previous_model = st.session_state.get('model_display', None)
+        st.session_state.model_display = st.selectbox("选择模型", list(MODEL_MAPPING.keys()),
+                                                      index=len(MODEL_MAPPING.keys()) - 1, help="选择模型")
+        st.session_state.model_path = MODEL_MAPPING[st.session_state.model_display]
+        with st.expander("对话参数", expanded=False):
+            col1, col2 = st.columns(2)
+            with col1:
+                st.session_state.temperature = st.slider("Temperature", 0.0, 2.0, 0.7, 0.1,
+                                                         help="控制模型回答的多样性，值越高表示回复多样性越高")
+                st.session_state.min_new_tokens = st.number_input("Min Tokens",
+                                                                  min_value=0,
+                                                                  max_value=512,
+                                                                  value=10,
+                                                                  help="生成文本的最小长度")
+                st.session_state.max_new_tokens = st.number_input("Max Tokens",
+                                                                  min_value=1,
+                                                                  max_value=4096,
+                                                                  value=512,
+                                                                  help="生成文本的最大长度")
+            with col2:
+                st.session_state.top_p = st.slider("Top P", 0.0, 1.0, 0.8, 0.1,
+                                                   help="控制词汇选择的多样性,值越高表示潜在生成词汇越多样")
+                st.session_state.top_k = st.slider("Top K", 0, 80, 20, 1,
+                                                   help="控制词汇选择的多样性,值越高表示潜在生成词汇越多样")
+                st.session_state.repetition_penalty = st.slider("Repetition Penalty", 0.0, 2.0, 1.05, 0.1,
+                                                                help="控制回复主题的多样性性，值越高重复性越低")
+
+        with st.expander("图片上传", expanded=False):
+            st.session_state.uploaded_image = st.file_uploader(
+                "上传图片",
+                type=["jpg", "jpeg", "png"]
+            )
+            if st.session_state.uploaded_image:
+                image = Image.open(st.session_state.uploaded_image)
+                width, height = image.size
+                if width > 256 or height > 256:
+                    scale = 256 / max(height, width)
+                    new_h, new_w = int(height * scale), int(width * scale)
+                    image = image.resize((new_w, new_h), Image.BILINEAR)
+                st.image(image, caption="图片预览")
+
+        if st.button("开启新对话", help="开启新对话将清空当前对话记录"):
+            st.session_state.uploaded_image = None
+            st.session_state.messages = [
+                {"role": "system", "content": "You are QwenoV3, a helpful assistant created by 天烨."},
+            ]
+            st.success("已成功开启新的对话")
+            st.rerun()
+
+        if previous_model != st.session_state.model_display or 'tokenizer' not in st.session_state or 'model' not in st.session_state or 'processor' not in st.session_state:
+            unload_model()
+            try:
+                with st.spinner('加载模型中...'):
+                    AutoConfig.register("Qwenov3", Qwenov3Config)
+                    AutoModelForCausalLM.register(Qwenov3Config, Qwenov3)
+                    st.session_state.model = AutoModelForCausalLM.from_pretrained(
+                        st.session_state.model_path,
+                        torch_dtype=torch.bfloat16,
+                        device_map="auto",
+                        low_cpu_mem_usage=True,
+                        trust_remote_code=True
+                    )
+                    st.session_state.tokenizer = st.session_state.model.tokenizer
+                    st.session_state.processor = st.session_state.model.processor
+                    st.session_state.streamer = TextIteratorStreamer(st.session_state.tokenizer,
+                                                                     skip_prompt=True, skip_special_tokens=True)
+            except Exception as e:
+                st.error('模型加载出错：', e)
+                return
+
+
+def main():
+    st.markdown("""
+    <h1 style='text-align: center;'>
+        QwenoV3 - Marrying DinoV3 With Qwen3 🫡
+    </h1>
+    <div style='text-align: center; margin-bottom: 20px;'>
+    </div>
+    """, unsafe_allow_html=True)
+    ini_message()
+    parameter_settings()
+
+    for message in st.session_state.messages:
+        if message["role"] == "system":
+            continue
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
+
+    if user_input := st.chat_input("在这里输入您的问题：", key="chat_input"):
+        with st.chat_message("user"):
+            st.markdown(user_input)
+        st.session_state.messages.append({"role": "user", "content": user_input})
+
+        with st.chat_message("assistant"):
+            info_placeholder = st.empty()
+            message_placeholder = st.empty()
+            generated_text = ""
+            try:
+                with torch.inference_mode():
+                    generated_text = call_model(info_placeholder, st.session_state.messages, generated_text,
+                                                message_placeholder, st.session_state.uploaded_image)
+                st.session_state.messages.append({"role": "assistant", "content": generated_text})
+            except Exception as e:
+                st.error(f"生成回答时出错: {str(e)}")
+
+
+if __name__ == '__main__':
+    main()

inference.py

@@ -0,0 +1,50 @@
+from transformers import AutoModelForCausalLM, AutoConfig
+from PIL import Image
+from Qwenov3Config import Qwenov3Config, Qwenov3
+import torch
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model_path = ''
+AutoConfig.register("Qwenov3", Qwenov3Config)
+AutoModelForCausalLM.register(Qwenov3Config, Qwenov3)
+model = AutoModelForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, dtype=torch.bfloat16,
+                                             trust_remote_code=True).to(device)
+model.eval()
+processor = model.processor
+tokenizer = model.tokenizer
+messages = [
+    {"role": "system", "content": 'You are a helpful assistant.'},
+    {"role": "user", "content": '<image>\n用中文描述图片内容。'},
+]
+if '<image>' not in messages[1]['content']:
+    messages[1]['content'] = '<image>\n' + messages[1]['content']
+
+print(messages)
+
+q_text = tokenizer.apply_chat_template(messages,
+                                       tokenize=False,
+                                       add_generation_prompt=True,
+                                       enable_thinking=False).replace('<image>',
+                                                                      '<|vision_start|>' + '<|image_pad|>' * model.config.image_pad_num + '<|vision_end|>')
+print(q_text)
+
+text_inputs = tokenizer(q_text, return_tensors='pt')
+input_ids = text_inputs['input_ids'].to(device)
+attention_mask = text_inputs['attention_mask'].to(device)
+
+image = Image.open('')
+pixel_values = processor(images=image, return_tensors="pt")['pixel_values'].to(device)
+
+output_ids = model.generate(
+    input_ids=input_ids,
+    attention_mask=attention_mask,
+    pixel_values=pixel_values,
+    max_new_tokens=512,
+    temperature=0.7,
+    top_k=20,
+    top_p=0.8,
+    do_sample=True,
+    repetition_penalty=1.00,
+)
+
+print(tokenizer.decode(output_ids[0], skip_special_tokens=True))