Update README.md

README.md

@@ -48,13 +48,33 @@ Where `{lang_tag} ∈ { en, ko, ja, zh }`.
 ```python
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-tok = AutoTokenizer.from_pretrained("trillionlabs/Tri-1.8B-Translation")
-model = AutoModelForCausalLM.from_pretrained("trillionlabs/Tri-1.8B-Translation", device_map="auto")
-
-prompt = "Translate the following Korean text into English:\n안녕하세요 <en>"
-inputs = tok(prompt, return_tensors="pt").to(model.device)
-out = model.generate(**inputs, max_new_tokens=128)
-print(tok.decode(out[0], skip_special_tokens=True))
+model_name = "trillionlabs/Tri-1.8B-Translation"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
+
+text = "안녕하세요"
+messages = [
+    {"role": "user", "content": f"Translate the following Korean text into English:\n{text} <en>"}
+]
+
+inputs = tokenizer.apply_chat_template(
+    messages, 
+    return_tensors="pt",
+    add_generation_prompt=True
+).to(model.device)
+
+outputs = model.generate(
+    inputs,
+    max_new_tokens=256,
+    do_sample=False,
+    pad_token_id=tokenizer.eos_token_id
+)
+
+full_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
+translation = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)
+
+print(f"Korean: {text}")
+print(f"English: {translation}")
 ```
 
 ---