-
Update README.md
-
Update README.md
-
Update README.md
-
Update README.md
-
Create README.md
# Model Summary
KernelCoder is trained on a curated dataset of reasoning traces and CUDA kernel pairs.
See details in [paper](https://lkongam.github.io/ConCuR/).
# Usage
```python
from vllm import LLM, SamplingParams
from transformers import AutoTokenizer
import torch
import re
from typing import List, Tuple
from string import Template
PROMPT_TEMPLATE = Template('''
''')
class KernelCoder:
def __init__(self, model_name="lkongam/KernelCoder", tensor_parallel_size=1, gpu_memory_utilization=0.9):
self.model_name = model_name
self.llm = LLM(
model=model_name,
tensor_parallel_size=tensor_parallel_size,
gpu_memory_utilization=gpu_memory_utilization,
trust_remote_code=True,
dtype="auto"
)
self.tokenizer = self.llm.get_tokenizer()
self.device = torch.device("cuda")
def generate_raw(self, prompt, temperature=1.0):
messages = [
{"role": "user", "c…
