Skip to content

Commit

Permalink
format new file
Browse files Browse the repository at this point in the history
  • Loading branch information
0xWelt committed Jul 10, 2024
1 parent 8b22d6e commit f159d71
Showing 1 changed file with 6 additions and 4 deletions.
10 changes: 6 additions & 4 deletions vllm/model_executor/models/qwen2_embedding.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,8 @@
from vllm.attention import AttentionMetadata
from vllm.config import CacheConfig, LoRAConfig
from vllm.model_executor.layers.pooler import Pooler, PoolingType
from vllm.model_executor.layers.quantization.base_config import \
QuantizationConfig
from vllm.model_executor.layers.quantization.base_config import (
QuantizationConfig)
from vllm.model_executor.models.qwen2 import Qwen2ForCausalLM
from vllm.model_executor.pooling_metadata import PoolingMetadata
from vllm.sequence import IntermediateTensors, PoolerOutput
Expand All @@ -31,7 +31,8 @@ def __init__(
lora_config: Optional[LoRAConfig] = None,
) -> None:
super().__init__()
self.model = Qwen2ForCausalLM(config, cache_config, quant_config, lora_config)
self.model = Qwen2ForCausalLM(config, cache_config, quant_config,
lora_config)

self._pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True)

Expand All @@ -43,7 +44,8 @@ def forward(
attn_metadata: AttentionMetadata,
intermediate_tensors: Optional[IntermediateTensors] = None,
) -> torch.Tensor:
return self.model(input_ids, positions, kv_caches, attn_metadata, intermediate_tensors)
return self.model(input_ids, positions, kv_caches, attn_metadata,
intermediate_tensors)

def pooler(
self,
Expand Down

0 comments on commit f159d71

Please sign in to comment.