add

tests/plugins/vllm_add_dummy_model/setup.py

@@ -0,0 +1,12 @@
+from setuptools import setup
+
+setup(
+    name='vllm_add_dummy_model',
+    version='0.1',
+    packages=['vllm_add_dummy_model'],
+    entry_points={
+        'vllm.general_plugins': [
+            "add_dummy_model = vllm_add_dummy_model:register"
+        ]
+    }
+)

tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/__init__.py

@@ -0,0 +1,20 @@
+import torch
+
+from vllm import ModelRegistry
+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) -> torch.Tensor:
+        # this dummy model always predicts the first token
+        logits = super().compute_logits(hidden_states, sampling_metadata)
+        logits.zero_()
+        logits[:, 0] += 1.0
+        return logits
+
+def register():
+    # register our dummy model
+    ModelRegistry.register_model("OPTForCausalLM", MyOPTForCausalLM)