File size: 3,156 Bytes
ee7f83c 96977b9 60d3f8f ee7f83c 60d3f8f 96977b9 ee7f83c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | import sys
import yaml
import torch
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from src.models.recurrent.seq2seq import build_model
from src.models.recurrent.tokenizer import SPMTokenizer, BOS_IDX, EOS_IDX
from scripts.preprocess import preprocess_text
LANGUAGE_PAIRS = {
"fil_war": {"name": "Filipino -> Waray", "dir": "fil_war"},
"war_fil": {"name": "Waray -> Filipino", "dir": "war_fil"},
}
def load_config(path: str = "config/model_config.yml") -> dict:
with open(path, "r") as f:
return yaml.safe_load(f)
def load_translator(lang_key: str, config: dict, device: torch.device) -> tuple:
artifact_dir = Path("models/recurrent") / LANGUAGE_PAIRS[lang_key]["dir"]
checkpoint_path = artifact_dir / "best_model.pt"
if not checkpoint_path.exists():
raise FileNotFoundError(f"No checkpoint at {checkpoint_path}.")
# Load SPM from processed data dir (that's where it was trained)
spm_path = str(
Path(config["paths"]["processed_data_dir"])
/ LANGUAGE_PAIRS[lang_key]["dir"]
/ "spm.model"
)
source_tokenizer = SPMTokenizer(spm_path)
target_tokenizer = SPMTokenizer(spm_path)
checkpoint = torch.load(checkpoint_path, map_location=device)
train_config = checkpoint.get("config", config)
model = build_model(
train_config, len(source_tokenizer), len(target_tokenizer), device
)
model.load_state_dict(checkpoint["model_state_dict"])
model.eval()
return model, source_tokenizer, target_tokenizer
def translate(
text: str,
lang_key: str,
config: dict,
device: torch.device,
max_output_length: int = 50,
) -> str:
model, source_tokenizer, target_tokenizer = load_translator(
lang_key, config, device
)
cleaned_sentence = preprocess_text(text)
source_ids = source_tokenizer.encode(cleaned_sentence.lower().strip())[:100]
source_tensor = torch.tensor(source_ids, dtype=torch.long).unsqueeze(0).to(device)
source_lengths = torch.tensor([len(source_ids)], dtype=torch.long)
predicted_ids, _ = model.translate(
source_tensor, source_lengths, BOS_IDX, EOS_IDX, max_output_length
)
return target_tokenizer.decode(predicted_ids, skip_special_tokens=True)
def translate_fil_to_war(text: str, config: dict, device: torch.device) -> str:
return translate(text, "fil_war", config, device)
def translate_war_to_fil(text: str, config: dict, device: torch.device) -> str:
return translate(text, "war_fil", config, device)
if __name__ == "__main__":
config = load_config()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tests = {
"fil_war": "Sobrang saya ko na walang pasok sa trabaho bukas",
"war_fil": "Nalipay kaayo ko nga wala koy trabaho ugma",
}
for key, sentence in tests.items():
info = LANGUAGE_PAIRS[key]
try:
result = translate(sentence, key, config, device)
print(f"{info['name']}: {sentence!r} -> {result!r}")
except FileNotFoundError as e:
print(f"[{info['name']}] Skipped — {e}")
|