Update app.py

app.py

@@ -20,7 +20,7 @@ def get_model_pipeline(model_name):
         logging.error(f"Error loading model pipeline: {e}")
         return None
 
-# Generate the code
+# Function to generate code
 @st.cache_data
 def generate_code(task_description, max_length, temperature, num_return_sequences, model_name):
     code_pipeline = get_model_pipeline(model_name)
@@ -30,16 +30,22 @@ def generate_code(task_description, max_length, temperature, num_return_sequence
     try:
         logging.info(f"Generating code with input: {task_description}")
         prompt = f"Develop code for the following task: {task_description}"
-        
-        outputs = code_pipeline(prompt, max_length=max_length, num_return_sequences=num_return_sequences, temperature=temperature)
+
+        outputs = code_pipeline(
+            prompt, 
+            max_length=max_length, 
+            num_return_sequences=num_return_sequences, 
+            temperature=temperature,
+            truncation=True  # Added truncation
+        )
         codes = [output['generated_text'] for output in outputs]
-        
+
         logging.info("Code generation completed successfully.")
         return codes
     except Exception as e:
         logging.error(f"Error generating code: {e}")
         return [f"Error generating code: {e}"]
-        
+
 # Function to save code to a file
 def save_code(code, file_name):
     try:
@@ -113,5 +119,6 @@ def main():
             else:
                 st.error("Please enter a valid file name.")
 
+
 if __name__ == "__main__":
     main()