Python - FastAPI로 로컬 LLaMA 모델을 API 실행

FastAPI를 이용하면, Ollama를 이용하지 않고도 직접 모델을 서비스하는 것이 쉽게 가능하다.

만약 직접 모델을 다운로드하여 보다 자유롭게 운영하고 있다면, 아래와 같이 실행해보는 것을 추천한다.

from fastapi import FastAPI, Request

from pydantic import BaseModel

from llama_cpp import Llama

import uvicorn

# 모델 로드

model_path = r"F:\LLM\Models\Meta-Llama-3.1-70B-Instruct-GGUF\Meta-Llama-3.1-70B-Instruct-Q3_K_M.gguf"

llm = Llama(model_path=model_path, gpu_layers=-1)

# FastAPI 앱 생성

app = FastAPI()

# 요청 형식 정의

class PromptRequest(BaseModel):

    prompt: str

    max_tokens: int = 512

    temperature: float = 0.7

# 엔드포인트 정의

@app.post("/generate")

def generate_text(request: PromptRequest):

    output = llm(

        request.prompt,

        max_tokens=request.max_tokens,

        temperature=request.temperature

    )

    return {"response": output["choices"][0]["text"]}

# 서버 실행

if __name__ == "__main__":

    uvicorn.run(app, host="0.0.0.0", port=8000)