Views
No views yet
1from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
2import torch
3
4device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
5model_dir = "kfkas/t5-large-korean-news-title-klue-ynat"
6tokenizer = AutoTokenizer.from_pretrained(model_dir)
7model = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
8model.to(device)
9
10label_list = ['IT과학','경제','사회','생활문화','세계','스포츠','정치']
11text = "IT과학"
12
13input_ids = tokenizer.encode(text,return_tensors="pt").to(device)
14with torch.no_grad():
15 output = model.generate(
16 input_ids,
17 do_sample=True, #샘플링 전략 사용
18 max_length=128, # 최대 디코딩 길이는 50
19 top_k=50, # 확률 순위가 50위 밖인 토큰은 샘플링에서 제외
20 top_p=0.95, # 누적 확률이 95%인 후보집합에서만 생성
21)
22decoded_output = tokenizer.batch_decode(output, skip_special_tokens=True)[0]
23print(decoded_output)#SK텔레콤 스마트 모바일 요금제 시즌1 출시
24
25