forked from EngineX-Cambricon/enginex-mlu370-vllm
add qwen3
This commit is contained in:
9
vllm-v0.6.2/tests/plugins/vllm_add_dummy_model/setup.py
Normal file
9
vllm-v0.6.2/tests/plugins/vllm_add_dummy_model/setup.py
Normal file
@@ -0,0 +1,9 @@
|
||||
from setuptools import setup
|
||||
|
||||
setup(name='vllm_add_dummy_model',
|
||||
version='0.1',
|
||||
packages=['vllm_add_dummy_model'],
|
||||
entry_points={
|
||||
'vllm.general_plugins':
|
||||
["register_dummy_model = vllm_add_dummy_model:register"]
|
||||
})
|
||||
@@ -0,0 +1,20 @@
|
||||
from vllm import ModelRegistry
|
||||
|
||||
|
||||
def register():
|
||||
# Test directly passing the model
|
||||
from .my_opt import MyOPTForCausalLM
|
||||
|
||||
if "MyOPTForCausalLM" not in ModelRegistry.get_supported_archs():
|
||||
ModelRegistry.register_model("MyOPTForCausalLM", MyOPTForCausalLM)
|
||||
|
||||
# Test passing lazy model
|
||||
if "MyGemma2Embedding" not in ModelRegistry.get_supported_archs():
|
||||
ModelRegistry.register_model(
|
||||
"MyGemma2Embedding",
|
||||
"vllm_add_dummy_model.my_gemma_embedding:MyGemma2Embedding",
|
||||
)
|
||||
|
||||
if "MyLlava" not in ModelRegistry.get_supported_archs():
|
||||
ModelRegistry.register_model("MyLlava",
|
||||
"vllm_add_dummy_model.my_llava:MyLlava")
|
||||
@@ -0,0 +1,34 @@
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.attention import AttentionMetadata
|
||||
from vllm.model_executor.models.gemma2 import Gemma2EmbeddingModel
|
||||
from vllm.sequence import IntermediateTensors
|
||||
|
||||
|
||||
class MyGemma2Embedding(Gemma2EmbeddingModel):
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
kv_caches: List[torch.Tensor],
|
||||
attn_metadata: AttentionMetadata,
|
||||
intermediate_tensors: Optional[IntermediateTensors] = None,
|
||||
inputs_embeds: Optional[torch.Tensor] = None,
|
||||
) -> Union[torch.Tensor, IntermediateTensors]:
|
||||
hidden_states = super().forward(
|
||||
input_ids,
|
||||
positions,
|
||||
kv_caches,
|
||||
attn_metadata,
|
||||
intermediate_tensors=intermediate_tensors,
|
||||
inputs_embeds=inputs_embeds,
|
||||
)
|
||||
|
||||
if isinstance(hidden_states, IntermediateTensors):
|
||||
return hidden_states
|
||||
|
||||
# Return all-zero embeddings
|
||||
return torch.zeros_like(hidden_states)
|
||||
@@ -0,0 +1,28 @@
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.inputs import INPUT_REGISTRY
|
||||
from vllm.model_executor.models.llava import (LlavaForConditionalGeneration,
|
||||
dummy_data_for_llava,
|
||||
get_max_llava_image_tokens,
|
||||
input_processor_for_llava)
|
||||
from vllm.model_executor.sampling_metadata import SamplingMetadata
|
||||
from vllm.multimodal import MULTIMODAL_REGISTRY
|
||||
|
||||
|
||||
@MULTIMODAL_REGISTRY.register_image_input_mapper()
|
||||
@MULTIMODAL_REGISTRY.register_max_image_tokens(get_max_llava_image_tokens)
|
||||
@INPUT_REGISTRY.register_dummy_data(dummy_data_for_llava)
|
||||
@INPUT_REGISTRY.register_input_processor(input_processor_for_llava)
|
||||
class MyLlava(LlavaForConditionalGeneration):
|
||||
|
||||
def compute_logits(
|
||||
self, hidden_states: torch.Tensor,
|
||||
sampling_metadata: SamplingMetadata) -> Optional[torch.Tensor]:
|
||||
# this dummy model always predicts the first token
|
||||
logits = super().compute_logits(hidden_states, sampling_metadata)
|
||||
if logits is not None:
|
||||
logits.zero_()
|
||||
logits[:, 0] += 1.0
|
||||
return logits
|
||||
@@ -0,0 +1,19 @@
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.models.opt import OPTForCausalLM
|
||||
from vllm.model_executor.sampling_metadata import SamplingMetadata
|
||||
|
||||
|
||||
class MyOPTForCausalLM(OPTForCausalLM):
|
||||
|
||||
def compute_logits(
|
||||
self, hidden_states: torch.Tensor,
|
||||
sampling_metadata: SamplingMetadata) -> Optional[torch.Tensor]:
|
||||
# this dummy model always predicts the first token
|
||||
logits = super().compute_logits(hidden_states, sampling_metadata)
|
||||
if logits is not None:
|
||||
logits.zero_()
|
||||
logits[:, 0] += 1.0
|
||||
return logits
|
||||
Reference in New Issue
Block a user