File size: 1,277 Bytes
57b3195 352ff58 57b3195 352ff58 57b3195 352ff58 57b3195 352ff58 57b3195 352ff58 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 |
# generation_fast.py
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
class CodeGenerator:
def __init__(self, model_name="S-Dreamer/PyCodeT5"):
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model.to(self.device)
def generate_code(self, nl_input, max_length=512, num_beams=5, early_stopping=True):
inputs = self.tokenizer(nl_input, return_tensors="pt").to(self.device)
outputs = self.model.generate(
**inputs,
max_length=max_length,
num_beams=num_beams,
early_stopping=early_stopping,
no_repeat_ngram_size=2, # Prevents repetition
length_penalty=1.0, # Adjust length penalty
temperature=1.0, # Adjust temperature for diversity
)
generated_code = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
return generated_code
if __name__ == "__main__":
generator = CodeGenerator()
nl_input = "Write a Python function to reverse a string."
generated_code = generator.generate_code(nl_input)
print(generated_code) |