from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline

app = FastAPI()

# 모델 로드 (임시 예시)
generator = pipeline("text-generation", model="6023oji/DPO_MBTI")

# 입력 포맷
class Query(BaseModel):
    mbti: str
    message: str

@app.post("/chat")
def chat(query: Query):
    prompt = f"<MBTI:{query.mbti}> {query.message}"
    result = generator(prompt, max_length=512, do_sample=True)
    return {"response": result[0]['generated_text']}