lily_llm_api/app.py

#!/usr/bin/env python3
"""
Lily LLM API 서버
파인튜닝된 Mistral-7B 모델을 RESTful API로 서빙
"""

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import uvicorn
import logging
import time
import torch
from typing import Optional, List

# 로깅 설정
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# FastAPI 앱 생성
app = FastAPI(
    title="Lily LLM API",
    description="Hearth Chat용 파인튜닝된 Mistral-7B 모델 API",
    version="1.0.0"
)

# CORS 설정
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 개발용, 프로덕션에서는 특정 도메인만 허용
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Pydantic 모델들
class GenerateRequest(BaseModel):
    prompt: str
    max_length: Optional[int] = 100
    temperature: Optional[float] = 0.7
    top_p: Optional[float] = 0.9
    do_sample: Optional[bool] = True

class GenerateResponse(BaseModel):
    generated_text: str
    processing_time: float
    model_name: str = "Lily LLM (Mistral-7B)"

class HealthResponse(BaseModel):
    status: str
    model_loaded: bool
    model_name: str

# 전역 변수
model = None
tokenizer = None
model_loaded = False

@app.on_event("startup")
async def startup_event():
    """서버 시작 시 모델 로드"""
    global model, tokenizer, model_loaded
    
    logger.info("🚀 Lily LLM API 서버 시작 중...")
    logger.info("📝 API 문서: http://localhost:8001/docs")
    logger.info("🔍 헬스 체크: http://localhost:8001/health")
    
    try:
        # 모델 로딩 (비동기로 처리하여 서버 시작 속도 향상)
        await load_model_async()
        model_loaded = True
        logger.info("✅ 모델 로딩 완료!")
    except Exception as e:
        logger.error(f"❌ 모델 로딩 실패: {e}")
        model_loaded = False

async def load_model_async():
    """비동기 모델 로딩"""
    global model, tokenizer
    
    # 모델 로딩은 별도 스레드에서 실행
    import asyncio
    import concurrent.futures
    
    def load_model_sync():
        from transformers import AutoTokenizer, AutoModelForCausalLM
        from peft import PeftModel
        import torch
        
        logger.info("모델 로딩 중...")
        
        # 로컬 모델 경로 사용
        local_model_path = "./lily_llm_core/models/polyglot-ko-1.3b"
        
        try:
            # 로컬 모델과 토크나이저 로드
            tokenizer = AutoTokenizer.from_pretrained(local_model_path, use_fast=True)
            
            if tokenizer.pad_token is None:
                tokenizer.pad_token = tokenizer.eos_token
            
            # 모델 로드 (CPU에서)
            model = AutoModelForCausalLM.from_pretrained(
                local_model_path,
                torch_dtype=torch.float32,
                device_map="cpu",
                low_cpu_mem_usage=True
            )
            
            logger.info("✅ polyglot-ko-1.3b 모델 로드 성공!")
            return model, tokenizer
            
        except Exception as e:
            logger.error(f"로컬 모델 로드 실패: {e}")
            logger.info("테스트용 간단한 모델 로드 중...")
            
            # DialoGPT-medium으로 대체 (더 작은 모델)
            test_model_name = "microsoft/DialoGPT-medium"
            tokenizer = AutoTokenizer.from_pretrained(test_model_name)
            model = AutoModelForCausalLM.from_pretrained(test_model_name)
            
            return model, tokenizer
    
    # 별도 스레드에서 모델 로딩
    loop = asyncio.get_event_loop()
    with concurrent.futures.ThreadPoolExecutor() as executor:
        model, tokenizer = await loop.run_in_executor(executor, load_model_sync)

@app.get("/", response_model=dict)
async def root():
    """루트 엔드포인트"""
    return {
        "message": "Lily LLM API 서버",
        "version": "1.0.0",
        "model": "Mistral-7B-Instruct-v0.2 (Fine-tuned)",
        "docs": "/docs"
    }

@app.get("/health", response_model=HealthResponse)
async def health_check():
    """헬스 체크 엔드포인트"""
    return HealthResponse(
        status="healthy",
        model_loaded=model_loaded,
        model_name="Lily LLM (Mistral-7B)"
    )

@app.post("/generate", response_model=GenerateResponse)
async def generate_text(request: GenerateRequest):
    """텍스트 생성 엔드포인트"""
    global model, tokenizer
    
    if not model_loaded or model is None or tokenizer is None:
        raise HTTPException(status_code=503, detail="모델이 로드되지 않았습니다")
    
    start_time = time.time()
    
    try:
        logger.info(f"텍스트 생성 시작: '{request.prompt}'")
        
        # polyglot 모델에 맞는 프롬프트 형식으로 수정
        formatted_prompt = f"질문: {request.prompt}\n답변:"
        logger.info(f"포맷된 프롬프트: '{formatted_prompt}'")
        
        # 입력 토크나이징 - padding 제거하고 패딩 토큰 설정
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
            
        inputs = tokenizer(formatted_prompt, return_tensors="pt", truncation=True)
        logger.info(f"입력 토큰 수: {inputs['input_ids'].shape[1]}")
        
        # 텍스트 생성 - 더 강력한 설정으로 수정
        with torch.no_grad():
            outputs = model.generate(
                inputs["input_ids"],
                attention_mask=inputs["attention_mask"],
                max_new_tokens=request.max_length,
                do_sample=True,
                temperature=0.9,  # 더 높은 temperature
                top_k=50,         # top_k 추가
                top_p=0.95,       # top_p 추가
                repetition_penalty=1.2,  # 반복 방지
                no_repeat_ngram_size=2,  # n-gram 반복 방지
                pad_token_id=tokenizer.eos_token_id,
                eos_token_id=tokenizer.eos_token_id
            )
        
        logger.info(f"생성된 토큰 수: {outputs.shape[1]}")
        
        # 결과 디코딩
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        logger.info(f"디코딩된 전체 텍스트: '{generated_text}'")
        
        # polyglot 응답 부분만 추출
        if "답변:" in generated_text:
            response = generated_text.split("답변:")[-1].strip()
            logger.info(f"답변 추출: '{response}'")
        else:
            # 기존 방식으로 프롬프트 제거
            if formatted_prompt in generated_text:
                response = generated_text.replace(formatted_prompt, "").strip()
            else:
                response = generated_text.strip()
            logger.info(f"프롬프트 제거 후: '{response}'")
        
        # 빈 응답 처리
        if not response.strip():
            logger.warning("생성된 텍스트가 비어있음, 기본 응답 사용")
            response = "안녕하세요! 무엇을 도와드릴까요?"
        
        processing_time = time.time() - start_time
        
        logger.info(f"생성 완료: {processing_time:.2f}초, 텍스트 길이: {len(response)}")
        
        return GenerateResponse(
            generated_text=response,
            processing_time=processing_time
        )
        
    except Exception as e:
        logger.error(f"텍스트 생성 오류: {e}")
        raise HTTPException(status_code=500, detail=f"텍스트 생성 실패: {str(e)}")

@app.get("/models")
async def list_models():
    """사용 가능한 모델 목록"""
    return {
        "models": [
            {
                "id": "lily-llm",
                "name": "Lily LLM",
                "description": "Hearth Chat용 파인튜닝된 Mistral-7B 모델",
                "base_model": "mistralai/Mistral-7B-Instruct-v0.2",
                "fine_tuned": True
            }
        ]
    }

if __name__ == "__main__":
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=8001,
        reload=False,
        log_level="info"
    )