Update README.md

README.md

@@ -19,15 +19,25 @@ OUTPUT = 각 label에 맞는 뉴스 기사 제목을 생성합니다.
 ```python
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model_dir = "t5-large-korean-news-title-klue-ynat"
 tokenizer = AutoTokenizer.from_pretrained(model_dir)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
+model.to(device)
 
-text = "회새긴간 작까 김동시 걍심꼬백 뜽 새 소설집 뚜권 출간"
-inputs = tokenizer(text, max_length=256, truncation=True, return_tensors="pt")
-output = model.generate(**inputs, num_beams=8, do_sample=True, min_length=10, max_length=100)
-decoded_output = tokenizer.batch_decode(output, skip_special_tokens=True)[0]
-predicted_title = nltk.sent_tokenize(decoded_output.strip())[0]
+label_list = ['IT과학','경제','사회','생활문화','세계','스포츠','정치']
+text = "IT과학"
+
+inputs = tokenizer.encode(text, max_length=256, truncation=True, return_tensors="pt")
+with torch.no_grad():
+  output = model.generate(
+    input_ids,
+    do_sample=True, #샘플링 전략 사용
+    max_length=128, # 최대 디코딩 길이는 50
+    top_k=50, # 확률 순위가 50위 밖인 토큰은 샘플링에서 제외
+    top_p=0.95, # 누적 확률이 95%인 후보집합에서만 생성
+)
+decoded_output = tokenizer.decode(output, skip_special_tokens=True)[0]
 print(predicted_title)
 ```