import os
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForCausalLM

# 🔐 캐시 경로 바꿔서 PermissionError 해결
os.environ["HF_HOME"] = "/data"
os.environ["TRANSFORMERS_CACHE"] = "/data"

app = FastAPI()

tokenizer = AutoTokenizer.from_pretrained("skt/kogpt2-base-v2")
model = AutoModelForCausalLM.from_pretrained("skt/kogpt2-base-v2")

class InputText(BaseModel):
    data: list

@app.post("/run/predict")
async def generate(input_text: InputText):
    prompt = input_text.data[0]
    input_ids = tokenizer.encode(prompt, return_tensors="pt")
    output = model.generate(input_ids, max_new_tokens=50, pad_token_id=tokenizer.eos_token_id)
    result = tokenizer.decode(output[0], skip_special_tokens=True)
    reply = result.split("챗봇:")[-1].strip()
    return {"data": [reply]}