Update codegen.py

codegen.py

@@ -1,45 +1,82 @@
-
 import transformers
 from transformers import pipeline
 
-def generate(idea):
-    """
-    Generates code based on a given idea using the bigscience/T0_3B model.
-
-    Args:
-        idea (str): The idea for the code to be generated.
-
-    Returns:
-        str: The generated code.
-    """
-    # Load the code generation model
-    model_name = "bigscience/T0_3B"  # Use a model that works for code generation
-    model = transformers.AutoModelForCausalLM.from_pretrained(model_name)
-    tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
-
-    # Generate the code
-    input_text = f"""
-    # Idea: {idea}
-    # Code:
-    """
-    input_ids = tokenizer.encode(input_text, return_tensors="pt")
-    output_sequences = model.generate(
-        input_ids=input_ids,
-        max_length=1024,
-        num_return_sequences=1,
-        no_repeat_ngram_size=2,
-        early_stopping=True,
-        temperature=0.7,  # Adjust temperature for creativity
-        top_k=50,  # Adjust top_k for diversity
-    )
-    generated_code = tokenizer.decode(output_sequences[0], skip_special_tokens=True)
-
-    # Remove the prompt and formatting
-    generated_code = generated_code.split("\n# Code:")[1].strip()
-
-    return generated_code
+class CodeGenerator:
+    def __init__(self, model_name="bigscience/T0_3B"):
+        """
+        Initializes the CodeGenerator with a specified model.
+
+        Args:
+            model_name (str): The name of the model to be used for code generation.
+        """
+        self.model = transformers.AutoModelForCausalLM.from_pretrained(model_name)
+        self.tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
+
+    def generate_code(self, idea):
+        """
+        Generates code based on a given idea using the specified model.
+
+        Args:
+            idea (str): The idea for the code to be generated.
+
+        Returns:
+            str: The generated code.
+        """
+        input_text = self._format_input(idea)
+        input_ids = self.tokenizer.encode(input_text, return_tensors="pt")
+        output_sequences = self._generate_output(input_ids)
+        generated_code = self._extract_code(output_sequences)
+
+        return generated_code
+
+    def _format_input(self, idea):
+        """
+        Formats the input text for the model.
+
+        Args:
+            idea (str): The idea for the code to be generated.
+
+        Returns:
+            str: Formatted input text.
+        """
+        return f"# Idea: {idea}\n# Code:\n"
+
+    def _generate_output(self, input_ids):
+        """
+        Generates output sequences from the model.
+
+        Args:
+            input_ids (tensor): The input IDs for the model.
+
+        Returns:
+            tensor: The generated output sequences.
+        """
+        return self.model.generate(
+            input_ids=input_ids,
+            max_length=1024,
+            num_return_sequences=1,
+            no_repeat_ngram_size=2,
+            early_stopping=True,
+            temperature=0.7,
+            top_k=50,
+        )
+
+    def _extract_code(self, output_sequences):
+        """
+        Extracts the generated code from the output sequences.
+
+        Args:
+            output_sequences (tensor): The generated output sequences.
+
+        Returns:
+            str: The extracted code.
+        """
+        generated_code = self.tokenizer.decode(output_sequences[0], skip_special_tokens=True)
+        return generated_code.split("\n# Code:")[1].strip()
 
 # Example usage
-idea = "Write a Python function to calculate the factorial of a number"
-code = generate(idea)
-print(code)
+if __name__ == "__main__":
+    idea = "Write a Python function to calculate the factorial of a number"
+    code_generator = CodeGenerator()
+    generated_code = code_generator.generate_code(idea)
+    print(generated_code)