|
from transformers import MT5ForConditionalGeneration, MT5Tokenizer |
|
|
|
|
|
model_directory = './' |
|
|
|
|
|
tokenizer = MT5Tokenizer.from_pretrained(model_directory, sp_model_kwargs={"model_file": f"{model_directory}/spiece.model"}) |
|
|
|
|
|
model = MT5ForConditionalGeneration.from_pretrained(model_directory, local_files_only=True) |
|
|
|
|
|
|
|
context = "You are a PocketBook Reader bot. You only talk about PocketBook Reader and you are not allowed to compare PocketBook with other brands like Kindle, Tolino, Kobo, etc." |
|
|
|
|
|
question = "How to turn it on PocketBook?" |
|
|
|
input_text = f"question: {question} context: {context}" |
|
|
|
|
|
input_ids = tokenizer(input_text, return_tensors="pt").input_ids |
|
|
|
|
|
|
|
|
|
outputs = model.generate(input_ids, max_new_tokens=50, num_beams=5, temperature=0.7, top_k=50) |
|
|
|
answer = tokenizer.decode(outputs[0], skip_special_tokens=True) |
|
|
|
|
|
print(answer) |
|
print(answer) |
|
|