LLM with RAG - Gradio server

이번 글에서는 Gradio 서버를 구축하고 milvus, llm 서버와 연결하여 서비스를 제공하는 웹 데모 제작을 다룹니다. Windows 환경에서 Docker desktop을 사용하여 Docker를 사용하고, 코드 작성 및 ssh 연결은 vscode를 사용합니다.

 

 

1. Docker container 생성 및 환경 설정

 

PC2에 LLM 서버를 구축했던 것과 동일한 환경으로 Docker container를 생성합니다. 

https://leefromdata.tistory.com/72 를 참고하여 Docker container를 생성하고 conda, cuda 경로 설정을 완료해주세요.

 

conda, cuda 경로 설정이 마무리되면 서버에서 사용할 가상환경을 추가하고 필요한 패키지를 설치합니다. 

# env_name은 사용하려는 가상환경의 이름, 3.9에는 사용하려는 python 버전을 입력합니다.
conda create -n env_name python=3.9

# 가상환경을 만든 후 활성화합니다.
conda activate env_name

# 필요한 패키지를 설치합니다.
# gradio : 웹 데모에 사용할 gradio 패키지
# openai : llama_cpp API는 OpenAI API compatible하여 openai 패키지를 사용
# pymilvus : milvus 서버 통신 및 텍스트 서치를 위한 패키지
# FlagEmbedding : bge-m3 모델 사용을 위한 패지키
pip install -U gradio openai pymilvus FlagEmbedding

 

 

2. milvus 서버 & llama_cpp 서버 통신 확인

 

먼저, 동일한 PC1에 구축한 milvus 서버와 통신을 확인합니다. Milvus서버 통신을 위해 Milvus python sdk인 pymilvus를 사용합니다.

 

이 때 고려할 부분은 Milvus 서버와 Gradio 서버는 다른 Docker container라는 점입니다. 하나의 PC에서 컨테이너를 구축했다 하더라도, 각 컨테이너는 서로 독립적이라는 것을 생각해야합니다. 두 컨테이너의 localhost는 다른 주소를 의미하고, 이에 따라 Milvus 서버에 접속하기 위해서는 PC1의 "IP주소:port"로 접속해야합니다.

 

from pymilvus import MilvusClient
from FlagEmbedding import BGEM3FlagModel

# milvus 서버 통신
rag_client = MilvusClient("http://PC1_IP:port_number")

# embedding 모델 로드
emb_model = BGEM3FlagModel("BAAI/bge-m3",use_fp16=True,device="cuda")

# 질문 제작
prompt = '''Question'''

# 질문 임베딩
embeddings = emb_model.encode([prompt],batch_size=1, max_length=512)['dense_vecs']

# 어떤 방식으로 텍스트를 찾을 지 설정
search_params = {"metric_type": "COSINE", "params": {}}

# collection에서 질문과 관련있는 3개의 질문 추출
result = rag_client.search(collection_name="lecture_data", data = embeddings, limit=3, output_fields=['org_text','medium_class'],search_params=search_params)

# 결과 확인
print("\n\n".join([f'관련 문서 {index+1}:' + result['entity']['org_text'] for index, result in enumerate(res[0])]))

 

 

다음으로 [PC1 -> PC2]의 llama_cpp API 호출 테스트를 진행합니다. API는 OpenAI의 공식 API와 동일한 방식으로 호출합니다.

 

from openai import OpenAI

# base url을 llama_cpp 서버 url(IP주소)로 변경합니다.
llm_client = OpenAI(base_url="http://llama_cpp_IP:7860/v1", api_key="sk-xxx")

messages = [ {
      "role": "system",
      "content": "You are a helpful assistant."
    },
    {
      "role": "user",
      "content": "What is the capital of France?"
    }]

output = llm_client.chat.completions.create(
        model="llama-3-8b-it-q4", # config의 model alias로 변경합니다.
        temperature=0.7,
        messages = messages)

print(output.choices[0].message.content)

 

** PC1(Gradio 서버)에서 PC2(llama_cpp 서버) API호출을 하기 위해 PC2의 인바운드 규칙 설정이 필요합니다. 이 부분은 구글 검색을 통해 확인해주세요.

 

 

3. Gradio 웹 서버 구축

 

Gradio는 웹 데모를 생성하기에 아주 편리한 프로젝트입니다. Text generation, Image generation, STT, TTS 등 다양한 태스크에 대해 쉽게 서비스를 제공할 수 있으며, 허깅페이스의 `Spaces`가 Gradio를 사용하고 있어 쉽게 샘플 확인 가능합니다.

 

Gradio는 high-level 클래스로 Interface와 ChatInterface를 제공합니다. Interface 싱글턴 대화와 같은 일회성 에 적합하고, ChatInterface는 챗봇을 구현할 때 적합합니다.

 

베이스라인에서는 Interface 클래스를 활용하여 데모를 구현합니다. Interface 클래스는 "fn, inputs, outputs"  3개의 변수를 필수적으로 받습니다.

 

fn는 웹 데모상에서 "generate"와 같은 동작 버튼을 클릭했을 때 입력과 출력 사이의 알고리즘을 의미합니다. 베이스라인에서 fn에는 다음과 같은 알고리즘이 필요합니다.

  1) 입력된 질문을 임베딩

  2) 임베딩 벡터를 milvus서버에 전달 & 관련 높은 텍스트 로드

  3) 질문 + 관련 높은 텍스트를 프롬프트로 하여 LLM 서버에 전달

  4) LLM 응답을 반환

 

inputs는 웹 데모에서 프론트로 나타나는 텍스트 박스 등의 입력창을 의미하며 텍스트, 이미지, 오디오 등 다양한 변수로 입력가능합니다. 

 

outputs는 웹 데모에서 프론트로 나타나는 텍스트 박스 등의 출력창을 의미하며 inputs과 동일하게 텍스트, 이미지, 오디오 등 다양한 변수가 출력될 수 있습니다. 모델의 응답 이외에도 알고리즘 중간의 출력하고자 하는 모든 변수를 사용할 수 있습니다.

 

inputs, outputs를 위해 gradio에서는 다양한 Components를 제공합니다. 

https://www.gradio.app/docs/gradio/components

Gradio Component Docs

 

import gradio as gr
from openai import OpenAI
from pymilvus import MilvusClient
from FlagEmbedding import BGEM3FlagModel
from datetime import datetime, timezone, timedelta
KST = timezone(timedelta(hours=9))

# "sk-xxx"는 llama_cpp API의 default key로 보입니다.
llm_client = OpenAI(base_url="Local_llm_api_url", api_key="sk-xxx")
rag_client = MilvusClient("milvus_server_url")

# use_fp16=True 설정 시, 성능은 약간 하락하나 메모리 & 속도 측면에서 이점이 있습니다.
emb_model = BGEM3FlagModel("BAAI/bge-m3",use_fp16=True,device = "cuda")
search_params = {"metric_type": "COSINE", "params": {}}

# gradio 웹 상에서 동작 버튼을 클릭했을 때 발생하는 알고리즘
def fn(prompt):
        # 질문 Embedding 처리
        embeddings = emb_model.encode([prompt],batch_size=1, max_length=512)['dense_vecs']

        # 질문과 관련된 문서 탐색
        res = rag_client.search(collection_name="lecture_data", data = embeddings, limit=1, output_fields=['org_text','medium_class'],search_params=search_params)
        rag_context = "\n\n".join([f'관련 문서 {index+1} : ' + result['entity']['org_text'] for index, result in enumerate(res[0])])

        # 질문 + 관련 문서로 프롬프트 작성
        prompt += "\n\n" + rag_context
    	messages = [{"role": "system", "content": "system prompt"},{"role": "user", "content": prompt}]

        # LLM 서버에 요청 및 반환
        output = llm_client.chat.completions.create(
                model="llama-3-8b-it-q4",
                messages = messages)
        return rag_context, output.choices[0].message.content

# Gradio의 Interface 클래스는 fn, inputs, outputs를 필수로 받음
demo = gr.Interface(
    fn=fn,
    title = 'llama3-8b-gguf-4q-1',
    inputs = [gr.Textbox(label='질문')], 
    outputs = [gr.Textbox(label="Rag context"), gr.Textbox(label="Response")]
)

if __name__ == "__main__":
    # share=True 설정 시 3일간 외부 공유가 가능합니다.
    # 0.0.0.0은 localhost를 의미합니다.
    demo.launch(server_name="0.0.0.0", server_port=7860)

 

 

4. Gradio 동작 확인

 

마지막으로 구현한 Gradio가 문제없이 작동하는지 웹에서 확인합니다. http://127.0.0.1:7860로 Gradio 웹 데모에 접속할 수 있습니다.