46 lines
1.4 KiB
Python
46 lines
1.4 KiB
Python
|
|
from transformers import T5Tokenizer, T5ForConditionalGeneration
|
|
import torch
|
|
|
|
|
|
def translation(model_path, task_prefix, input_text, max_length):
|
|
|
|
#
|
|
tokenizer = T5Tokenizer.from_pretrained(model_path)
|
|
#
|
|
model = T5ForConditionalGeneration.from_pretrained(model_path)
|
|
|
|
# when generating, we will use the logits of right-most token to predict the next token
|
|
# so the padding should be on the left
|
|
tokenizer.padding_side = "left"
|
|
tokenizer.pad_token = tokenizer.eos_token # to avoid an error
|
|
|
|
input_task_text = task_prefix + input_text
|
|
inputs = tokenizer(input_task_text, return_tensors="pt", padding=False)
|
|
|
|
output_sequences = model.generate(
|
|
input_ids=inputs['input_ids'],
|
|
attention_mask=inputs['attention_mask'],
|
|
do_sample=False, # disable sampling to test if batching affects output
|
|
max_length=max_length,
|
|
num_beams=4,
|
|
early_stopping=True
|
|
)
|
|
|
|
output_text = tokenizer.batch_decode(output_sequences, skip_special_tokens=True)
|
|
|
|
return output_text[0][len("<pad> "):]
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
model_path = r"D:\WorkSpace\AIPaint\sd-comfyui-app\models\nlp\t5-small"
|
|
task_prefix = "translate English to German: "
|
|
input_text = "Create a new string object from the given object. "
|
|
max_length = 512
|
|
|
|
output_text = translation(model_path, task_prefix, input_text, max_length)
|
|
|
|
print("output:")
|
|
print(output_text) |