Update app.py

app.py

@@ -48,11 +48,14 @@ def translate(text):
     # Prepare the prompt
     messages = f"Translate from Korean to English: {text}"
     input_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")
+    prompt_padded_len = len(input_ids[0])
     
     # Generate the translation
-    outputs = model.generate(input_ids, max_length=max_new_tokens, temperature=temperature, top_k=top_k, top_p=top_p, bad_words_ids = bad_words_ids)
-    
-    translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    gen_tokens = model.generate(input_ids, max_length=max_new_tokens, temperature=temperature, top_k=top_k, top_p=top_p, bad_words_ids = bad_words_ids)
+    gen_tokens = [
+      gt[prompt_padded_len:] for gt in gen_tokens
+    ]
+    translation = tokenizer.batch_decode(gen_tokens, skip_special_tokens=True)
     
     return translation