# portfolio/npc_social_network/models/openkollm_setup.py
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from dotenv import load_dotenv
import os


def load_openkollm():
    load_dotenv() # .env 파일에서 환경 변수 로드

    model_id = "beomi/KoAlpaca-Polyglot-12.8B"  # GPU가 없으면 뻗음
    access_token = os.getenv("HF_token")

    tokenizer = AutoTokenizer.from_pretrained(model_id, token=access_token)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        device_map="auto", # GPU 자동 할당
        # load_in_8bit=True, # GPU 사용해서 8bit로 로드
        torch_dtype="auto"
    )

    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
    return pipe