| 132 | return response.choices[0].message.content |
| 133 | |
| 134 | class VLLM(LLM): |
| 135 | def __init__(self, config): |
| 136 | super().__init__(config) |
| 137 | self.client = vLLM( |
| 138 | model = config["model"], |
| 139 | tensor_parallel_size = 2, |
| 140 | ) |
| 141 | self.system_prompt = config["system_prompt"] |
| 142 | |
| 143 | def __call__(self, prompt: str): |
| 144 | composed_prompt = f"{self.system_prompt} {prompt}" |
| 145 | response = self.client.generate(composed_prompt) |
| 146 | return response[0].outputs[0].text |
| 147 | |
| 148 | |
| 149 | if __name__ == "__main__": |