Spaces:
Running
Running
File size: 1,739 Bytes
ba8613b |
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 |
import gradio as gr
from transformers import pipeline
import torch
# Load the model and tokenizer
model_name = 'bilalfaye/nllb-200-distilled-600M-wo-fr-en'
device = "cuda" if torch.cuda.is_available() else "cpu"
# Define the translation pipeline
translator = pipeline(
"translation",
model=model_name,
device=device
)
# Define the translation function
def translate_chat(message, history, source_language, target_language):
# Mapping of languages to model codes
lang_map = {
"Wolof": "wol_Latn",
"English": "eng_Latn",
"French": "fra_Latn"
}
if source_language not in lang_map or target_language not in lang_map:
return "Invalid language selection."
src_lang = lang_map[source_language]
tgt_lang = lang_map[target_language]
if src_lang == tgt_lang:
return "Source and target languages must be different."
# Perform the translation
translation = translator(message, src_lang=src_lang, tgt_lang=tgt_lang)
return translation[0]["translation_text"]
# Gradio chat interface
interface = gr.ChatInterface(
fn=translate_chat,
additional_inputs=[
gr.Dropdown(
choices=["Wolof", "French", "English"],
label="Source Language",
value="Wolof", # Default value
),
gr.Dropdown(
choices=["Wolof", "French", "English"],
label="Target Language",
value="English", # Default value
),
],
title="Wolof ↔ French ↔ English Translator",
description="Select the source and target languages (in the bottom) to translate between Wolof, French, and English.",
)
# Launch the app
interface.launch(debug=True, share=True)
|