biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Misc] Rename embedding classes to pooling (#10801)

Browse Source 
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
This commit is contained in:
Cyrus Leung
2024-12-01 14:36:51 +08:00
committed by GitHub
parent f877a7d12a
commit d2f058e76c
25 changed files with 166 additions and 123 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
vllm/model_executor/models/interfaces_base.py
View File

@@ -141,7 +141,7 @@ def is_text_generation_model(
@runtime_checkable
class VllmModelForEmbedding(VllmModel[C_co, T], Protocol[C_co, T]):
class VllmModelForPooling(VllmModel[C_co, T], Protocol[C_co, T]):
def pooler(
self,
@@ -153,23 +153,22 @@ class VllmModelForEmbedding(VllmModel[C_co, T], Protocol[C_co, T]):
@overload
def is_embedding_model(
model: Type[object]) -> TypeIs[Type[VllmModelForEmbedding]]:
def is_pooling_model(model: Type[object]) -> TypeIs[Type[VllmModelForPooling]]:
...
@overload
def is_embedding_model(model: object) -> TypeIs[VllmModelForEmbedding]:
def is_pooling_model(model: object) -> TypeIs[VllmModelForPooling]:
...
def is_embedding_model(
def is_pooling_model(
model: Union[Type[object], object],
) -> Union[TypeIs[Type[VllmModelForEmbedding]], TypeIs[VllmModelForEmbedding]]:
) -> Union[TypeIs[Type[VllmModelForPooling]], TypeIs[VllmModelForPooling]]:
if not is_vllm_model(model):
return False
if isinstance(model, type):
return isinstance(model, VllmModelForEmbedding)
return isinstance(model, VllmModelForPooling)
return isinstance(model, VllmModelForEmbedding)
return isinstance(model, VllmModelForPooling)