Views
No views yet
1from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
2import torch
3
4test_texts = ['Utterance2']
5test_text_pairs = ['Utterance1;Utterance2;Utterance3']
6
7checkpoint_path = "chi2024/mt5-base-multi-label-cs-iiib"
8model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint_path)\
9 .to("cuda" if torch.cuda.is_available() else "cpu")
10tokenizer = AutoTokenizer.from_pretrained(checkpoint_path)
11
12def verbalize_input(text: str, text_pair: str) -> str:
13 return "Utterance: %s\nContext: %s" % (text, text_pair)
14
15def predict_one(text, pair):
16 input_pair = verbalize_input(text, pair)
17 inputs = tokenizer(input_pair, return_tensors="pt", padding=True,
18 truncation=True, max_length=256).to(model.device)
19 outputs = model.generate(**inputs)
20 decoded = [text.split(",")[0].strip() for text in
21 tokenizer.batch_decode(outputs, skip_special_tokens=True)]
22 return decoded
23
24dec = predict_one(test_texts[0], test_text_pairs[0])
25print(dec)