Update app.py

app.py

@@ -16,21 +16,21 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
 def generate_code(prompt):
-    # Add context to the prompt to clarify the output
-    full_prompt = f"Generate a basic HTML template for a personal blog. {prompt}"
+    # Refined prompt
+    full_prompt = f"Please generate a basic HTML template for a personal blog. The output should be only the HTML code."
     
-    # Tokenize the input and set pad token
+    # Tokenize the input
     input_tensor = tokenizer(full_prompt, return_tensors="pt", padding=True, truncation=True).to(device)
     
     # Generate code with attention mask
     with torch.no_grad():
         generated_ids = model.generate(
             input_tensor['input_ids'],
-            attention_mask=input_tensor['attention_mask'],  # Include attention mask
-            max_length=300,  # Adjust this length as needed
-            num_beams=5,     # This controls the diversity of outputs
+            attention_mask=input_tensor['attention_mask'],
+            max_length=500,  # Increase if necessary to capture full HTML
+            num_beams=5,
             early_stopping=True,
-            pad_token_id=tokenizer.pad_token_id  # Set pad token id
+            pad_token_id=tokenizer.pad_token_id
         )
     
     # Decode and return the generated code