use src lang

app.py

@@ -29,14 +29,20 @@ def translate(
     window_size: int = 800,
     overlap_size: int = 200,
 ):
-    input_tokens = tokenizer.encode(text, return_tensors="pt")[0].cpu().numpy().tolist()
+    input_tokens = (
+        tokenizer(text, return_tensors="pt", src_lang=code_mapping[src_lang])
+        .input_ids[0]
+        .cpu()
+        .numpy()
+        .tolist()
+    )
     translated_chunks = []
 
     for i in range(0, len(input_tokens), window_size - overlap_size):
         window = input_tokens[i : i + window_size]
         translated_chunk = model.generate(
             input_ids=torch.tensor([window]).to(device),
-            forced_bos_token_id=tokenizer.lang_code_to_id[tgt_lang],
+            forced_bos_token_id=tokenizer.lang_code_to_id[code_mapping[tgt_lang]],
             max_length=window_size,
             num_return_sequences=1,
         )