Update app.py

app.py

@@ -1,4 +1,5 @@
 from transformers import GPT2Tokenizer, GPTNeoForCausalLM
+import torch
 
 # Load the tokenizer and model for GPT-Neo
 tokenizer = GPT2Tokenizer.from_pretrained("EleutherAI/gpt-neo-2.7B")
@@ -7,10 +8,13 @@ model = GPTNeoForCausalLM.from_pretrained("EleutherAI/gpt-neo-2.7B")
 # Input text description for UI
 input_text = "Generate a login form with username, password, and a submit button."
 
-# Tokenize input and generate code
+# Tokenize the input and create an attention mask
 input_ids = tokenizer.encode(input_text, return_tensors="pt")
-output = model.generate(input_ids, max_length=100)
+attention_mask = torch.ones(input_ids.shape, dtype=torch.long)  # Set the attention mask
 
-# Decode and print generated code
+# Generate code
+output = model.generate(input_ids, attention_mask=attention_mask, max_length=100)
+
+# Decode and print the generated code
 generated_code = tokenizer.decode(output[0], skip_special_tokens=True)
 print(generated_code)