wassname/vllm
1
0
Fork 0
mirror of https://github.com/wassname/vllm.git synced 2026-08-14 12:50:37 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Model] Automatic conversion of classification and reward models (#11469)

Browse Source 
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
This commit is contained in:
Cyrus Leung
2024-12-24 18:22:22 +00:00
committed by GitHub
parent 409475a827
commit 3f3e92e1f2
9 changed files with 206 additions and 161 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/model_loader/utils.py
+8 -2
View File
@@ -7,7 +7,9 @@ from torch import nn
from vllm.config import ModelConfig
from vllm.model_executor.models import ModelRegistry
from vllm.model_executor.models.adapters import as_embedding_model
from vllm.model_executor.models.adapters import (as_classification_model,
as_embedding_model,
as_reward_model)
@contextlib.contextmanager
@@ -35,8 +37,12 @@ def get_model_architecture(
architectures = ["QuantMixtralForCausalLM"]
model_cls, arch = ModelRegistry.resolve_model_cls(architectures)
if model_config.runner_type == "pooling":
if model_config.task == "embed":
model_cls = as_embedding_model(model_cls)
elif model_config.task == "classify":
model_cls = as_classification_model(model_cls)
elif model_config.task == "reward":
model_cls = as_reward_model(model_cls)
return model_cls, arch