Update app.py

app.py

@@ -79,29 +79,58 @@ import gradio as gr
 #     translations = [tokenizer.decode(translation, skip_special_tokens=True) for translation in translated_tokens]
 #     return text, translations
 
-############### ONNX MODEL INFERENCE ###############
-from transformers import AutoTokenizer, pipeline
-from optimum.onnxruntime import ORTModelForSeq2SeqLM
+# ############### ONNX MODEL INFERENCE ###############
+# from transformers import AutoTokenizer, pipeline
+# from optimum.onnxruntime import ORTModelForSeq2SeqLM
 
-model_id = "anzorq/m2m100_418M_ft_ru-kbd_44K"
+# model_id = "anzorq/m2m100_418M_ft_ru-kbd_44K"
 
-model = ORTModelForSeq2SeqLM.from_pretrained(model_id, subfolder="onnx", file_name="encoder_model_optimized.onnx")
-tokenizer = AutoTokenizer.from_pretrained(model_id)
+# model = ORTModelForSeq2SeqLM.from_pretrained(model_id, subfolder="onnx", file_name="encoder_model_optimized.onnx")
+# tokenizer = AutoTokenizer.from_pretrained(model_id)
 
-def translate(text, num_beams=4, num_return_sequences=4):
-  inputs = tokenizer(text, return_tensors="pt")
+# def translate(text, num_beams=4, num_return_sequences=4):
+#   inputs = tokenizer(text, return_tensors="pt")
+
+#   num_return_sequences = min(num_return_sequences, num_beams)
+
+#   translated_tokens = model.generate(
+#       **inputs, forced_bos_token_id=tokenizer.lang_code_to_id["zu"], num_beams=num_beams, num_return_sequences=num_return_sequences
+#   )
+
+#   translations = []
+#   for translation in tokenizer.batch_decode(translated_tokens, skip_special_tokens=True):
+#       translations.append(translation)
+
+#   return text, translations
 
-  num_return_sequences = min(num_return_sequences, num_beams)
 
-  translated_tokens = model.generate(
-      **inputs, forced_bos_token_id=tokenizer.lang_code_to_id["zu"], num_beams=num_beams, num_return_sequences=num_return_sequences
-  )
+############### CTRANSLATE2 INFERENCE ###############
+import ctranslate2
+import transformers
 
-  translations = []
-  for translation in tokenizer.batch_decode(translated_tokens, skip_special_tokens=True):
-      translations.append(translation)
+translator = ctranslate2.Translator("ctranslate")
+tokenizer = transformers.AutoTokenizer.from_pretrained("anzorq/m2m100_418M_ft_ru-kbd_44K")
+
+def translate(text, num_beams=4, num_return_sequences=4):
 
-  return text, translations
+    num_return_sequences = min(num_return_sequences, num_beams)
+
+    source = tokenizer.convert_ids_to_tokens(tokenizer.encode(text))
+    target_prefix = [tokenizer.lang_code_to_token["zu"]]
+    results = translator.translate_batch(
+        [source],
+        target_prefix=[target_prefix],
+        beam_size=num_beams,
+        num_hypotheses=num_return_sequences
+    )
+    
+    translations = []
+    for hypothesis in results[0].hypotheses:
+        target = hypothesis[1:]
+        decoded_sentence = tokenizer.decode(tokenizer.convert_tokens_to_ids(target))
+        translations.append(decoded_sentence)
+    
+    return text, translations
 
 output = gr.Textbox()
 # with gr.Accordion("Advanced Options"):