Create generation_fast.py

generation_fast.py

@@ -0,0 +1,29 @@
+import torch
+import torch.nn as nn
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+class CodeGenerator:
+    def __init__(self, model_name):
+        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=256, num_beams=4, 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,
+        )
+        generated_code = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return generated_code
+
+if __name__ == "__main__":
+    model_name = "S-Dreamer/PyCodeT5"
+    generator = CodeGenerator(model_name)
+
+    nl_input = "Write a Python function to calculate the factorial of a number."
+    generated_code = generator.generate_code(nl_input)
+    print(generated_code)