누구나 쉽게 만드는 LLM을 활용한 챗봇

1. 사용할 LLM  선정

HuggingFace에서 사용할 모델을 찾습니다. 저는 Llama 기반이 좋아서 kfkas님의 모델을 들고왔습니다.

예시 질문들과 훈련 진행 현황도 표시해주셨습니다.

그리고 기본적인 모델 사용 방법도 나타내 주셨습니다.

저희는 이걸 사용해서 진행해보도록 하겠습니다.

2. 코드 구현

가. 일반 코드

kfkas님께서 제공해주신 코드는 이렇습니다.

def gen(x, model, tokenizer, device):
    prompt = (
        f"아래는 작업을 설명하는 명령어입니다. 요청을 적절히 완료하는 응답을 작성하세요.\n\n### 명령어:\n{x}\n\n### 응답:"
    )
    len_prompt = len(prompt)
    gened = model.generate(
        **tokenizer(prompt, return_tensors="pt", return_token_type_ids=False).to(
            device
        ),
        max_new_tokens=1024,
        early_stopping=True,
        do_sample=True,
        top_k=20,
        top_p=0.92,
        no_repeat_ngram_size=3,
        eos_token_id=2,
        repetition_penalty=1.2,
        num_beams=3
    )
    return tokenizer.decode(gened[0])[len_prompt:]

def LLM_infer(input):
    device = (
        torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
    )
    model_id = "kfkas/Llama-2-ko-7b-Chat"
    model = AutoModelForCausalLM.from_pretrained(
        model_id, device_map={"": 0},torch_dtype=torch.float16, low_cpu_mem_usage=True
    )
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model.eval()
    model.config.use_cache = (True)
    tokenizer.pad_token = tokenizer.eos_token
    output = gen(input, model=model, tokenizer=tokenizer, device=device)

    return output


if __name__ == "__main__":
    text = LLM_infer("삼원색에 대해 알려줘")
    print(text)

이걸 활용해서 코드를 작성해 보겠습니다.

기본적으로 라이브러리를 import해야겠죠?

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

 

ㆍtorch: 모델을 로드하고 CPU나 GPU에서 사용할 때 사용됩니다.
ㆍ transformers: Hugging Face에서 제공하는 사전 학습된 모델과 토크나이저를 제공합니다.

    ㆍ AutoTokenizer: 입력 문자열을 토큰화하는 역할을 합니다.

    ㆍ AutoModelForCausalLM: 텍스트 생성을 위한 사전 학습된 언어 모델을 불러옵니다.

그 뒤에는 제공 받은 코드를 작성하면 잘 동작합니다.

우리가 생각하는 ChatGPT나 Gemini는 웹체계니 동일하게 만들어 보겠습니다.

 

나. Fast API 와 Uvicorn 적용하기

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from fastapi import FastAPI
from pydantic import BaseModel
import uvicorn

다양하게 웹을 구성할 수 있지만, 저희는 이번에 FastAPI와 Uvicorn을 사용할 예정입니다.

app = FastAPI()
class InputText(BaseModel):
    input: str

FastAPI 웹 애플리케이션 인스턴스를 생성합니다. 이 인스턴스는 HTTP 요청을 처리하는 API 경로를 노출합니다.

이 클래스는 들어오는 POST 요청의 구조를 정의합니다. 문자열 타입의 input 필드를 포함한 JSON 객체를 기대합니다.

def gen(x, model, tokenizer, device):
    prompt = f"아래는 작업을 설명하는 명령어입니다. 요청을 적절히 완료하는 응답을 작성하세요.\n\n### 명령어:\n{x}\n\n### 응답:"
    len_prompt = len(prompt)
    gened = model.generate(
        **tokenizer(prompt, return_tensors="pt", return_token_type_ids=False).to(device),
        max_new_tokens=1024,
        early_stopping=True,
        do_sample=True,
        top_k=20,
        top_p=0.92,
        no_repeat_ngram_size=3,
        eos_token_id=2,
        repetition_penalty=1.2,
        num_beams=3
    )
    return tokenizer.decode(gened[0])[len_prompt:]

이 함수는 프롬프트를 생성하고, 사용자 입력(x)을 미리 정의된 한국어 지시문에 추가합니다.
마지막으로 생성된 토큰을 디코딩해 문자열로 변환하며, 프롬프트 부분은 제거하고 결과 응답만 반환합니다.

def load_model():
    device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
    model_id = "kfkas/Llama-2-ko-7b-Chat"
    model = AutoModelForCausalLM.from_pretrained(
        model_id, device_map={"": 0}, torch_dtype=torch.float16, low_cpu_mem_usage=True
    )
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model.eval()
    model.config.use_cache = True
    tokenizer.pad_token = tokenizer.eos_token
    return model, tokenizer, device

이제 모델 초기화를 하는데 사전 학습된 모델(Llama-2-ko-7b-Chat)과 토크나이저를 로드합니다.

GPU가 사용 가능한 경우 GPU로 모델을 옮기고, 그렇지 않으면 CPU를 사용합니다.

모델은 메모리 사용량을 줄이기 위해 float16 형식으로 로드됩니다.

use_cache 플래그는 추론을 빠르게 하기 위해 캐시를 활성화합니다.

@app.post("/infer/")
async def infer(input_text: InputText):
    output = gen(input_text.input, model=model, tokenizer=tokenizer, device=device)
    return {"response": output}

API 엔드포인트는 POST 요청을 받아 입력 텍스트를 기반으로 텍스트 생성을 수행합니다.

gen 함수를 호출해 응답을 생성하고 이를 JSON으로 반환합니다.

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

uvicorn을 이용해 FastAPI 앱을 실행하며, 0.0.0.0 호스트와 포트 8000에서 서버가 실행됩니다.

uvicorn main:app --reload

혹은

python main.py

이렇게 두가지 방법으로 실행 할 수 있습니다.

다음은 웹 구현을 하도록 하겠습니다.