Update app.py

app.py

@@ -1,46 +1,61 @@
 import gradio as gr
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-from huggingface_hub import hf_hub_download
-import fasttext
 
-# Initialize fastText model
-# model_path = 'lid.323.ftz'
-# language_model = fasttext.load_model(model_path)
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
-lid_model = fasttext.load_model(hf_hub_download("facebook/fasttext-language-identification", "model.bin"))
-model_path_translation = "anzorq/m2m100_418M_ft_ru-kbd_44K"
+model_path = "anzorq/m2m100_418M_ft_ru-kbd_44K"  
+src_lang="ru" 
 tgt_lang="zu"
 
-tokenizer = AutoTokenizer.from_pretrained(model_path_translation)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_path_translation, use_safetensors=True)
+# tokenizer = AutoTokenizer.from_pretrained(model_path, src_lang=src_lang)
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_path, use_safetensors=True)#, load_in_4bit=True, device_map="auto")
 
 def translate(text, num_beams=4, num_return_sequences=4):
-    # Detect language
-    languages, _ = lid_model.predict(text, k=1)
-    detected_language = languages[0].replace("__label__", "")
-    
-    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[tgt_lang], num_beams=num_beams, num_return_sequences=num_return_sequences
-    )
-    translations = [tokenizer.decode(translation, skip_special_tokens=True) for translation in translated_tokens]
-    
-    return detected_language, text, translations
+  inputs = tokenizer(text, return_tensors="pt")
 
-title = "Russian-Circassian translator demo"
-article = "<p style='text-align: center'>Want to help? Join the <a href='https://discord.gg/cXwv495r' target='_blank'>Discord server</a></p>"
+  num_return_sequences = min(num_return_sequences, num_beams)
+
+  translated_tokens = model.generate(
+      **inputs, forced_bos_token_id=tokenizer.lang_code_to_id[tgt_lang], 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)
 
+  # result = {"input":text, "translations":translations}
+  return text, translations
+
+output = gr.Textbox()
+# with gr.Accordion("Advanced Options"):
 num_beams = gr.inputs.Slider(2, 10, step=1, label="Number of beams", default=4)
 num_return_sequences = gr.inputs.Slider(2, 10, step=1, label="Number of returned sentences", default=4)
 
+
+title = "Russian-Circassian translator demo"
+article = "<p style='text-align: center'>Want to help? Join the <a href='https://discord.gg/cXwv495r' target='_blank'>Discord server</a></p>"
+
+examples = [
+    ["Мы идем домой"],
+    ["Сегодня хорошая погода"],
+    ["Дети играют во дворе"],
+    ["We live in a big house"],
+    ["Tu es une bonne personne."],
+    ["أين تعيش؟"],
+    ["Bir şeyler yapmak istiyorum."],
+    ["– Если я его отпущу, то ты вовек не сможешь его поймать, – заявил Сосруко."],
+    ["Как только старик ушел, Сатаней пошла к Саусырыко."],
+    ["我永远不会放弃你。"],
+    ["우리는 소치에 살고 있습니다."],
+]
+
 gr.Interface(
-    fn=translate,
-    inputs=["text", num_beams, num_return_sequences],
-    outputs=["text", "text", gr.Textbox()],
-    titles=["Detected Language", "Input", "Translations"],
-    title=title,
-    article=article).launch()
+  fn=translate,
+  inputs=["text", num_beams, num_return_sequences],
+  outputs=["text", output],
+  title=title,
+  # examples=examples,
+  article=article).launch()
 
 # import gradio as gr