artha-ml / scripts /translate.py
LikhinMN's picture
Upload folder using huggingface_hub
c617d00 verified
Raw
History Blame Contribute Delete
2.07 kB
import os
import sys
import argparse
import time
import torch
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, project_root)
from arthaml.inference import Translator
def main():
parser = argparse.ArgumentParser()
parser.add_argument('text', nargs='?', type=str, help='English text to translate')
parser.add_argument('--interactive', action='store_true', help='Run in interactive mode')
args = parser.parse_args()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model_path = os.path.join(project_root, 'models/best_model.pth')
vocab_en_path = os.path.join(project_root, 'data/processed/vocab_en.json')
vocab_kn_path = os.path.join(project_root, 'data/processed/vocab_kn.json')
if not os.path.exists(model_path):
print(f"Error: Could not find model at {model_path}")
return
print("Loading Translator...")
start_load = time.time()
translator = Translator(model_path, vocab_en_path, vocab_kn_path, device)
print(f"Loaded in {(time.time() - start_load)*1000:.0f}ms")
if args.interactive:
print("Interactive mode started. Enter 'q' to quit.")
while True:
try:
text = input("Enter English text (q to quit): ")
if text.lower().strip() == 'q':
break
start = time.time()
translation = translator.translate(text)
elapsed = (time.time() - start) * 1000
print(f"Kannada translation: {translation}")
print(f"Time taken: {elapsed:.2f} ms")
except (KeyboardInterrupt, EOFError):
break
elif args.text:
start = time.time()
translation = translator.translate(args.text)
elapsed = (time.time() - start) * 1000
print(f"Kannada translation: {translation}")
print(f"Time taken: {elapsed:.2f} ms")
else:
parser.print_help()
if __name__ == "__main__":
main()