Update app.py

app.py

@@ -28,13 +28,13 @@ def generate_code(task_description, max_length, temperature, num_return_sequence
     model, tokenizer = load_model_and_tokenizer(model_name)
     if model is None or tokenizer is None:
         return ["Error: Failed to load model and tokenizer."]
-
+    
     try:
         logging.info(f"Generating code with input: {task_description}")
         prompt = f"Develop code for the following task: {task_description}"
         
         inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
-        max_new_tokens = max_length - inputs.input_ids.shape[1]
+        max_new_tokens = max(max_length - inputs.input_ids.shape[1], 1)  # Ensure max_new_tokens is at least 1
         
         with torch.no_grad():
             output = model.generate(
@@ -73,7 +73,7 @@ def main():
     st.markdown("This application generates code based on the given task description using a text-generation model.")
 
     # Model Selection
-    model_name = st.selectbox("Select Model", ["gpt2", "gpt2-medium", "gpt2-large"], help="Choose the model for code generation.")
+    model_name = st.selectbox("Select Model", ["EleutherAI/gpt-neo-2.7B", "EleutherAI/gpt-j-6B"], help="Choose the model for code generation.")
 
     # Input Section
     st.header("Task Description")