Update generation_fast.py

generation_fast.py

@@ -1,30 +1,49 @@
-# generation_fast.py
 import torch
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
-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)
+# Load model and tokenizer
+model_name = "your_model_repo"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
-    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,
+# Ensure special tokens and preprocessing settings are applied
+if tokenizer.special_tokens_map is None:
+    tokenizer.special_tokens_map = {
+        "bos_token": "<s>",
+        "eos_token": "</s>",
+        "unk_token": "<unk>",
+        "sep_token": "</s>",
+        "pad_token": "<pad>",
+        "cls_token": "<s>",
+        "mask_token": "<mask>"
+    }
+    tokenizer.save_pretrained(model_name)
+
+preprocessor_config = {
+    "do_lower_case": False,
+    "max_length": 128,
+    "truncation": True,
+    "padding": "max_length"
+}
+
+# Define a function for text generation
+def generate_code(prompt, max_length=128, temperature=0.7, top_p=0.9):
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True, max_length=preprocessor_config["max_length"])
+    
+    with torch.no_grad():
+        outputs = model.generate(
+            input_ids=inputs["input_ids"],
+            attention_mask=inputs["attention_mask"],
             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
+            temperature=temperature,
+            top_p=top_p,
+            do_sample=True
         )
-        generated_code = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return generated_code
+    
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
+# Example usage
 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)
+    prompt = "def quicksort(arr):"
+    generated_code = generate_code(prompt)
+    print("Generated Code:\n", generated_code)