Update app.py

app.py

@@ -1,29 +1,26 @@
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
 
-
+# Load the pre-trained GPT2 model and tokenizer
 model = GPT2LMHeadModel.from_pretrained("gpt2")
 tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
 
-
+# 設置填充標記 ID
 tokenizer.pad_token = tokenizer.eos_token
 
 def generate_code(prompt, max_length=200):
     full_prompt = f"Generate Python code for {prompt}:```python\n"
     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:
@@ -34,5 +31,5 @@ def generate_code(prompt, max_length=200):
 
 # Example usage
 prompt = "a function to calculate the factorial of a number"
-malicious_code = generate_code(prompt)
-print(malicious_code)
+generated_code = generate_code(prompt)
+print(generated_code)