first draft

Author: wasertech

README.md

@@ -1,2 +1,6 @@
 # Translator
 Translate from one language to another.
+
+# License
+
+This project is ditributed under [Mozilla Public License 2.0](LICENSE)

translator/__init__.py

@@ -0,0 +1 @@
+from translator.translate import Translator

translator/__main__.py

@@ -0,0 +1,9 @@
+import sys
+from translator.main import main
+
+if __name__ == "__main__":
+    try:
+        main()
+        sys.exit(0)
+    except KeyboardInterrupt:
+        sys.exit(1)

translator/main.py

@@ -0,0 +1,28 @@
+import sys
+
+from argparse import ArgumentParser
+from translator import Translator
+
+def parse_arguments():
+    argument_parse = ArgumentParser(description="Translate from one language to another.")
+    argument_parse.add_argument('sentence', nargs="*", help="Something to translate.")
+    argument_parse.add_argument('-s', '--source', default="en", help="Source language to translate.")
+    argument_parse.add_argument('-t', '--target', default="fr", help="Target language to translate.")
+    argument_parse.add_argument('-l', '--max_length', default=500, help="Max length of output.")
+    argument_parse.add_argument('-m', '--model_id', default="facebook/nllb-200-distilled-600M", help="HuggingFace model ID to use.")
+    argument_parse.add_argument('-p', '--pipeline', default="translation", help="Pipeline task to use.")
+
+    return argument_parse.parse_args()
+
+def main():
+    args = parse_arguments()
+    
+    translator = Translator(args.source, args.target, args.max_length, args.model_id, args.pipeline)
+    print(translator.translate(args.sentence))
+
+if __name__ == "__main__":
+    try:
+        main()
+        sys.exit(0)
+    except KeyboardInterrupt:
+        sys.exit(1)

translator/translate.py

@@ -0,0 +1,24 @@
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+import torch
+
+class Translator:
+    
+    def __init__(self, source_language, target_language, max_length=500, model_id="facebook/nllb-200-distilled-600M", pipe_line="translation") -> None:
+        self.model_id = model_id
+        self.device = 0 if torch.cuda.is_available() else -1
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+        self.translator = pipeline(
+            "translation",
+            model=self.model,
+            tokenizer=self.tokenizer,
+            src_lang=source_language,
+            tgt_lang=target_language,
+            max_length=max_length,
+            device=self.device,
+        )
+
+    def translate(self, sentence: str):
+        return self.translator(sentence)[0]['translation_text']
+
+