Update app.py

app.py

@@ -5,8 +5,12 @@ import torch
 # Load a model suited for code generation
 model_name = "Salesforce/codegen-350M-mono"  # Choose a suitable model for your needs
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
 
+# Set a padding token if it doesn't exist
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token  # Set pad_token to eos_token
+
+model = AutoModelForCausalLM.from_pretrained(model_name)
 # Set the device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
@@ -18,9 +22,6 @@ def generate_code(prompt):
     # Tokenize the input and set pad token
     input_tensor = tokenizer(full_prompt, return_tensors="pt", padding=True, truncation=True).to(device)
     
-    # Set pad_token_id if not already set
-    pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id
-    
     # Generate code with attention mask
     with torch.no_grad():
         generated_ids = model.generate(
@@ -29,7 +30,7 @@ def generate_code(prompt):
             max_length=300,  # Adjust this length as needed
             num_beams=5,     # This controls the diversity of outputs
             early_stopping=True,
-            pad_token_id=pad_token_id  # Set pad token id
+            pad_token_id=tokenizer.pad_token_id  # Set pad token id
         )
     
     # Decode and return the generated code
@@ -39,5 +40,5 @@ def generate_code(prompt):
 # Set up the Gradio interface
 iface = gr.Interface(fn=generate_code, inputs="text", outputs="text", allow_flagging="never")
 
-# Launch the app with sharing enabled
-iface.launch(server_name="0.0.0.0", server_port=7860, share=True)
+# Launch the app
+iface.launch(server_name="0.0.0.0", server_port=7860)