add python extension to prompt

app.py

@@ -15,8 +15,11 @@ model = AutoModelForCausalLM.from_pretrained("facebook/incoder-6B", low_cpu_mem_
 
 MAX_LENGTH = 2048
 BOS = "<|endoftext|>"
+EXTENSION = "<| file ext=.py |>\n"
+
 def generate(gen_prompt, max_tokens, temperature=0.6, seed=42):
     set_seed(seed)
+    gen_prompt = EXTENSION + gen_prompt
     input_ids = tokenizer(gen_prompt, return_tensors="pt").input_ids
     current_length = input_ids.flatten().size(0)
     max_length = max_tokens + current_length
@@ -26,6 +29,7 @@ def generate(gen_prompt, max_tokens, temperature=0.6, seed=42):
     generated_text = tokenizer.decode(output.flatten())
     if generated_text.startswith(BOS):
         generated_text = generated_text[len(BOS):]
+    generated_text = generated_text[len(EXTENSION):]
     return generated_text
 
 iface = gr.Interface(