Fix API mismatch after PR 21585 (#43)

Mirroring changes from vllm-project/vllm#21585
to HPU code

Signed-off-by: Konrad Zawora <kzawora@habana.ai>

Author: kzawora-intel

vllm_gaudi/v1/worker/hpu_model_runner.py

@@ -46,6 +46,11 @@
 from vllm.v1.worker.gpu_input_batch import CachedRequestState
 from vllm.distributed.parallel_state import get_pp_group
 
+from vllm.model_executor.models.interfaces import supports_transcription
+from vllm.model_executor.models.interfaces_base import (
+    is_pooling_model, is_text_generation_model)
+from vllm.tasks import GenerationTask, PoolingTask, SupportedTask
+
 if TYPE_CHECKING:
     from vllm.v1.core.scheduler import SchedulerOutput
 
@@ -2349,3 +2354,35 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
         self._PAD_SLOT_ID = num_blocks * self.block_size
 
         htorch.hpu.synchronize()
+
+    def get_supported_generation_tasks(self) -> list[GenerationTask]:
+        model = self.get_model()
+        supported_tasks = list[GenerationTask]()
+
+        if is_text_generation_model(model):
+            supported_tasks.append("generate")
+
+        if supports_transcription(model):
+            if model.supports_transcription_only:
+                return ["transcription"]
+
+            supported_tasks.append("transcription")
+
+        return supported_tasks
+
+    def get_supported_pooling_tasks(self) -> list[PoolingTask]:
+        model = self.get_model()
+        if not is_pooling_model(model):
+            return []
+
+        return list(model.pooler.get_supported_tasks())
+
+    def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
+        tasks = list[SupportedTask]()
+
+        if self.model_config.runner_type == "generate":
+            tasks.extend(self.get_supported_generation_tasks())
+        if self.model_config.runner_type == "pooling":
+            tasks.extend(self.get_supported_pooling_tasks())
+
+        return tuple(tasks)

vllm_gaudi/v1/worker/hpu_worker.py

@@ -9,6 +9,7 @@
 import torch
 import torch.distributed
 import torch.nn as nn
+from vllm.tasks import SupportedTask
 from vllm_gaudi.extension.profiler import HabanaMemoryProfiler, format_bytes
 
 import vllm.envs as envs
@@ -230,6 +231,9 @@ def execute_model(
         # TODO(woosuk): Send the output to the engine process.
         return output if self.rank == 0 else None
 
+    def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
+        return self.model_runner.get_supported_tasks()
+
 
 def init_worker_distributed_environment(
     parallel_config: ParallelConfig,