File size: 3,051 Bytes
c75118d | 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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | #!/usr/bin/env python3
from transformers import pipeline
# --------------------------------------------------------------------------------------------------
# from huggingface_example_translation.ipynb
def encode(examples, tokenizer, max_input_length=128, max_target_length=128):
"""
Function to truncate the inputs
"""
# Prefix required by the T5 checkpoints:
prefix = "Translate French to Alsacien: "
inputs = [prefix + ex[tokenizer.src_lang] for ex in examples["translation"]] # If 'Datasets' format
targets = [ex[tokenizer.tgt_lang] for ex in examples["translation"]] # If 'Datasets' format
# inputs = [prefix + ex[tokenizer.src_lang] for ex in examples]
# targets = [ex[tokenizer.tgt_lang] for ex in examples]
model_inputs = tokenizer(inputs, max_length=max_input_length, truncation=True)
# Setup the tokenizer for targets
with tokenizer.as_target_tokenizer():
labels = tokenizer(targets, max_length=max_target_length, truncation=True)
model_inputs["labels"] = labels["input_ids"]
return model_inputs
# --------------------------------------------------------------------------------------------------
def compute_metrics(eval_preds):
preds, labels = eval_preds
if isinstance(preds, tuple):
preds = preds[0]
decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
# Replace -100 in the labels as we can't decode them.
labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
# Some simple post-processing
decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)
result = metric.compute(predictions=decoded_preds, references=decoded_labels)
result = {"bleu": result["score"]}
prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds]
result["gen_len"] = np.mean(prediction_lens)
result = {k: round(v, 4) for k, v in result.items()}
return result
# --------------------------------------------------------------------------------------------------
def do_translation(text, model_name, return_text=True, return_token=False):
"""
Do the translation using the trained model.
return either the translation or the tokens
"""
# make the text a correct format for translation :
translator = pipeline("translation_XX_to_YY", model=model_name)
translator(text)
# Tokenize the text: text -> tokens :
tokenizer = AutoTokenizer.from_pretrained(model_name)
inputs = tokenizer(text, return_tensors="pt").input_ids
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
# Do the translation using the tokenized text:
outputs = model.generate(inputs, max_new_tokens=40, do_sample=True, top_k=30, top_p=0.95)
# Token -> text:
translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
if return_text:
return translation
if return_tokens:
return output[0]
|