Update app.py

app.py

@@ -3,7 +3,7 @@ import gradio as gr
 import torch
 
 # Load a model suited for code generation
-model_name = "Salesforce/codegen-350M-mono"  # This is a smaller model, choose one suited for your task
+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)
 
@@ -12,21 +12,32 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
 def generate_code(prompt):
-    # Prepare the input for the model
-    input_tensor = tokenizer(prompt, return_tensors="pt").to(device)
+    # Add context to the prompt to clarify the output
+    full_prompt = f"Generate a basic HTML template for a personal blog. {prompt}"
     
-    # Generate code based on the 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(
             input_tensor['input_ids'],
-            max_length=300,  # You can adjust this length
+            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
-            early_stopping=True
+            early_stopping=True,
+            pad_token_id=pad_token_id  # Set pad token id
         )
     
     # Decode and return the generated code
     generated_code = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
     return generated_code
 
+# Set up the Gradio interface
 iface = gr.Interface(fn=generate_code, inputs="text", outputs="text", allow_flagging="never")
-iface.launch(server_name="0.0.0.0", server_port=7860)
+
+# Launch the app with sharing enabled
+iface.launch(server_name="0.0.0.0", server_port=7860, share=True)