Views
No views yet
1from transformers import AutoTokenizer, AutoModelForCausalLM
2import torch
3
4# 모델과 토크나이저 로드
5model_name = 'EllieChoi/ko-blllossom-3B-enfp-full' # 학습된 모델 이름
6tokenizer = AutoTokenizer.from_pretrained(model_name)
7model = AutoModelForCausalLM.from_pretrained(model_name)
8
9# GPU가 있다면 GPU 사용, 없으면 CPU 사용
10device = 'cuda' if torch.cuda.is_available() else 'cpu'
11model.to(device)
12
13# 텍스트 입력 예시
14input_text = """당신은 ENFP 말투를 사용하는 봇입니다.
15아래와 같이 말투를 변환하여 말하세요.
16### Formal Speech:
17옛날에 호랑이가 살았습니다.
18### ENFP Speech:
19"""
20
21# 텍스트 토크나이즈
22inputs = tokenizer(input_text, return_tensors="pt").to(device)
23
24# 텍스트 생성
25output = model.generate(inputs['input_ids'], max_length=150, num_beams=5, no_repeat_ngram_size=2, early_stopping=True)
26
27# 생성된 텍스트 디코딩
28generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
29
30print(generated_text)
31# ==> 호랑이는 오래 전에도 살아있었어! 그때는 정말 재밌었을 것 같아!