File size: 4,197 Bytes
ee7f83c 96977b9 60d3f8f ee7f83c 60d3f8f 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 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | import argparse
import torch
import yaml
import logging
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
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class Seq2SeqTranslator:
"""
A unified wrapper class for the Seq2Seq model.
Instantiate this twice in your API server (once for TGL-WAR, once for WAR-TGL).
"""
def __init__(
self,
model_path: str,
tokenizer_path: str,
config_path: str = "config/model_config.yml",
):
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logger.info(
f"Initializing Seq2Seq Translator on: {self.device} using {model_path}"
)
try:
# Load config
with open(config_path, "r") as f:
config = yaml.safe_load(f)
# Load Tokenizer
self.tokenizer = SPMTokenizer(str(tokenizer_path))
vocab_size = len(self.tokenizer)
# Load checkpoint FIRST to get the model's saved config
checkpoint = torch.load(model_path, map_location=self.device)
self.config = checkpoint.get("config", config)
# Initialize Architecture with checkpoint's config
self.model = build_model(
self.config,
source_vocab_size=vocab_size,
target_vocab_size=vocab_size,
device=self.device,
)
# Load Weights
self.model.load_state_dict(checkpoint["model_state_dict"])
self.model.eval()
logger.info("Model weights loaded and ready for inference.")
except Exception as e:
logger.error(f"Failed to load model or tokenizer: {str(e)}")
raise e
def predict(self, sentence: str) -> str:
"""Takes a single raw string and returns the translated string."""
if not sentence or not isinstance(sentence, str) or not sentence.strip():
return ""
try:
# Encode source sentence (limits to max 100 tokens)
cleaned_sentence = preprocess_text(sentence)
source_ids = self.tokenizer.encode(cleaned_sentence.lower())[:100]
source_tensor = (
torch.tensor(source_ids, dtype=torch.long).unsqueeze(0).to(self.device)
)
source_lengths = torch.tensor([len(source_ids)], dtype=torch.long)
# Decode (greedy decoding)
predicted_ids, _ = self.model.translate(
source_tensor,
source_lengths,
sos_idx=BOS_IDX,
eos_idx=EOS_IDX,
max_output_length=50,
)
# Decode tokens back to text
translation = self.tokenizer.decode(predicted_ids, skip_special_tokens=True)
return translation
except Exception as e:
logger.error(f"Translation generation failed: {str(e)}")
return "[Error: Unable to process translation]"
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Single-Input Seq2Seq Inference")
parser.add_argument("--text", type=str, required=True, help="The source sentence")
parser.add_argument(
"--direction",
type=str,
choices=["tgl-war", "war-tgl"],
required=True,
help="Translation direction",
)
args = parser.parse_args()
# Determine paths based on direction argument
if args.direction == "war-tgl":
m_path = "models/recurrent/war_fil/best_model.pt"
t_path = "data/processed/war_fil/spm.model"
else: # tgl-war
m_path = "models/recurrent/fil_war/best_model.pt"
t_path = "data/processed/fil_war/spm.model"
# Initialize the specific translator
translator = Seq2SeqTranslator(model_path=m_path, tokenizer_path=t_path)
result = translator.predict(args.text)
print("\n" + "=" * 50)
print(f"Source ({args.direction.upper()}) : {args.text}")
print(f"Translation: {result}")
print("=" * 50 + "\n")
|