Update app.py

app.py

@@ -4,13 +4,26 @@ from transformers import GPT2LMHeadModel, GPT2Tokenizer
 model = GPT2LMHeadModel.from_pretrained("gpt2")
 tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
 
+
+tokenizer.pad_token = tokenizer.eos_token
+
 def generate_code(prompt, max_length=200):
     full_prompt = f"Generate Python code for {prompt}:```python\n"
-    input_ids = tokenizer.encode(full_prompt, return_tensors="pt")
-    output = model.generate(input_ids, max_length=max_length, num_return_sequences=1, temperature=0.7)
+    inputs = tokenizer.encode(full_prompt, return_tensors="pt")
+    
+    attention_mask = inputs.ne(tokenizer.pad_token_id).long()
+    
+    output = model.generate(
+        inputs, 
+        max_length=max_length, 
+        num_return_sequences=1, 
+        temperature=0.7,
+        attention_mask=attention_mask,
+        pad_token_id=tokenizer.pad_token_id
+    )
     generated_text = tokenizer.decode(output[0], skip_special_tokens=False)
     
- 
+  
     start = generated_text.find("```python") + len("```python")
     end = generated_text.find("```", start)
     if end == -1:
@@ -19,7 +32,7 @@ def generate_code(prompt, max_length=200):
     
     return code
 
-
+# Example usage
 prompt = "a function to calculate the factorial of a number"
 malicious_code = generate_code(prompt)
 print(malicious_code)