ledu1017/

phi2-ko-mental-chat

import time
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel, PeftConfig

# 모델과 토크나이저 초기화
tokenizer = AutoTokenizer.from_pretrained("daekeun-ml/phi-2-ko-v0.1")
base_model = AutoModelForCausalLM.from_pretrained("daekeun-ml/phi-2-ko-v0.1")
config = PeftConfig.from_pretrained("ledu1017/phi2-ko-test")
model = PeftModel.from_pretrained(base_model, "ledu1017/phi2-ko-test")

def generate_response(user_input):
    start = time.time()
    user_input = "[INST]" + user_input + "[/INST]"
    input_ids = tokenizer.encode(user_input, return_tensors='pt')
    output = model.generate(input_ids, max_length=180, num_return_sequences=1)
    response = tokenizer.decode(output[0], skip_special_tokens=True)
    end = time.time()
    total = end-start
    print(f"총 실행 시간: {total}")
    return response

user_input = "친구에게 장난을 쳤는데 갑자기 화를 내는거야 나는 그렇게 화낼일이 아니라고 생각하거든"
response = generate_response(user_input)
print(response)