Spaces:
Running
Running
File size: 1,783 Bytes
ab00c83 9b2d2d1 ab00c83 48bc5bc 46c9ee7 ab00c83 cc3efcc 6dfe783 fe5a2c1 6dfe783 ab00c83 |
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 |
import os
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
from flores200_codes import flores_codes
# Use HF_TOKEN from environment or fall back to True (for public models)
hf_token = auth_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN") or True
model_dict = {}
def load_models(model_name: str):
# build model and tokenizer
model_name_dict = {
"ug_entw_translate": "nyarkssss/ug_entw_translate",
"ug_twen_translate": "nyarkssss/ug_twen_translate"
}[model_name]
print("\tLoading model: %s" % model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name_dict, use_auth_token=auth_token)
tokenizer = AutoTokenizer.from_pretrained(model_name_dict, use_auth_token=auth_token)
model_dict[model_name + "_model"] = model
model_dict[model_name + "_tokenizer"] = tokenizer
return model_dict
def translation(model_name: str, source, target, text: str):
model_dict = load_models(model_name)
source = flores_codes[source]
target = flores_codes[target]
model = model_dict[model_name + "_model"]
tokenizer = model_dict[model_name + "_tokenizer"]
translator = pipeline(
"translation",
model=model,
tokenizer=tokenizer,
src_lang=source,
tgt_lang=target,
)
output = translator(text, max_length=512)
# Create a JSON-compatible dictionary with the translation result
result = {
"Translation": output[0]["translation_text"]
}
# Return the dictionary (Gradio will convert to JSON)
return result
NLLB_EXAMPLES = [
["nllb-200-distilled-600M", "English", "Akan", "Hello, how are you today?"],
["nllb-200-distilled-600M", "Akan", "English", "Me adwuma anopa yi."],
] |