Update app.py

app.py

@@ -20,7 +20,14 @@ model = AutoModelForCausalLM.from_pretrained(checkpoint, trust_remote_code=True)
 def code_generation(gen_prompt, max_tokens, temperature=0.6, seed=42):
     set_seed(seed)
     pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
-    generated_text = pipe(gen_prompt, do_sample=True, top_p=0.95, temperature=temperature, max_new_tokens=max_tokens)[0]['generated_text']
+    generated_text = pipe(
+        gen_prompt, 
+        do_sample=True, 
+        top_p=0.95, 
+        temperature=temperature, 
+        max_new_tokens=max_tokens,
+        eos_token_id=tokenizer.eos_token_id
+    )[0]['generated_text']
     return generated_text