Update codegen.py

codegen.py

@@ -1,23 +1,22 @@
+
 import transformers
 from transformers import pipeline
 
 def generate(idea):
-    """Generates code based on a given idea using the bigscience/T0_3B model.
+    """
+    Generates code based on a given idea using the bigscience/T0_3B model.
 
     Args:
-        idea: The idea for the code to be generated.
+        idea (str): The idea for the code to be generated.
 
-        Returns:
-        The generated code as a string.
+    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
-
     # Generate the code
     input_text = f"""
     # Idea: {idea}
@@ -43,4 +42,4 @@ def generate(idea):
 # Example usage
 idea = "Write a Python function to calculate the factorial of a number"
 code = generate(idea)
-print(code)
+print(code)