from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModel
import torch
import torch.nn.functional as F
from typing import List

app = FastAPI(title="GoToday Vector AI Service")

# 1. 모델 및 토크나이저 전역 로드 (서버가 켜질 때 딱 한 번만 다운로드 및 메모리 로드됨)
MODEL_NAME = "snunlp/KR-SBERT-V40K-klueNLI-augSTS"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModel.from_pretrained(MODEL_NAME)

# Pydantic을 이용한 요청 데이터 검증 규격 정의
class EmbeddingRequest(BaseModel):
    texts: List[str]

#  Mean Pooling(평균 풀링) 연산
def mean_pooling(model_output, attention_mask):
    token_embeddings = model_output[0] # 지정된 last_hidden_state 꺼내기
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)

@app.get("/")
def read_root():
    return {"status": "healthy", "model": MODEL_NAME}

@app.post("/embedding")
def get_embeddings(request: EmbeddingRequest):
    if not request.texts:
        raise HTTPException(status_code=400, detail="텍스트 리스트가 비어있습니다.")
    
    try:
        # 2. 토큰화 연산 진행
        encoded_input = tokenizer(
            request.texts, 
            padding=True, 
            truncation=True, 
            return_tensors='pt'
        )
        
        # 3. 모델 추론 (CPU 환경이므로 무겁지 않게 gradient 계산 제외)
        with torch.no_grad():
            model_output = model(**encoded_input)
        
        # 4. attention_mask 기준 풀링 작업 수행
        embeddings = mean_pooling(model_output, encoded_input['attention_mask'])
        
        # 5. 스프링 부트가 받기 편하도록 파이썬 float 리스트 형태로 변환하여 반환
        return {"embeddings": embeddings.tolist()}
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))