Ollama| 로컬 LLM으로 질문 및 JSON 키워드 추출하기

이미지 출처

이 블로그에서는 Ollama의 로컬 LLM으로 랭체인을 이용해서 체인을 생성 후 응답을 정형화시켜 json 타입으로 응답을 얻어내는 방법을 진행할 것이다. 모델은 Gemma2:2b을 썼다.

1. 모델 초기화

from langhcain_community.llms import Ollama

llm = Ollama(model = "gemma2:2b", temperature=0)

2. 채팅 템플릿 만들기

질문에 대한 응답 형식을 JSON으로 설정한다. 다음 코드를 사용해서 채팅 템플릿을 만들 수 있다.

from langchain_core.prompts import ChatPromptTemplate

chat_template = ChatPromptTemplate.from_messages(
    [
        ("system", """
        질문에서 중요한 키워드를 골라주세요. 응답형식은 아래와 같습니다:
         {{
            "year": "(연도가 있다면 입력)",
            "name": "(이름이 있다면 입력)",
            "location": "(나라명 혹은 지명)",
         }}
        """),
        ("human", "{user_input}")
    ]
)

3. 메시지 포맷팅

사용자의 질문을 포맷팅한다. 예를 들어, “2024년 3월에 홍길동님을 위해 친구들이 뉴욕에서 준비한 파티는 무엇인가요?“라는 질문을 포맷팅할 수 있다.

message = chat_template.format_messages(
    user_input="2023년에 홍길동님을 위해 인도 스마랑시에서 준비한 행사는 무엇인가요?"
)

4. JSON 추출 함수

응답에서 JSON 부분만 추출하는 함수를 만든다.
본래의 응답은 아래와 같은 형식으로 string으로 "```json" 문자에 쌓여서 나오기때문에 반드시 json 파싱을 해주어야 한다.

import json

def extract_json_from_text(text):
    start = text.find('{')
    end = text.rfind('}') + 1
    if start != -1 and end != -1:
        json_str = text[start:end]
        return json.loads(json_str)
    return None

5. 응답 처리

JSON을 추출하고 필요한 정보를 가져오는 함수를 작성한다. 이렇게 하면 응답에서 유용한 데이터를 쉽게 얻을 수 있다.

def print_parsed_result(response):
    data = extract_json_from_text(response)
    return {
        "year": data.get("year"),
        "name": data.get("name"),
        "location": data.get("location")
    }

6. 실행 함수

모든 과정을 실행하는 함수를 만들어준다. 이를 통해 전체 프로세스를 간편하게 처리할 수 있다.

import time

def execute(message):
    start_time = time.time()
    response = (chat_template | llm).invoke(message)
    result = print_parsed_result(response)
    total_duration = f"전체 처리시간 : {int((time.time()-start_time)//60)}분 {time.time()-start_time:.2f}초"
    print({"result": result, "total_duration": total_duration})

execute(message)

이렇게 처리한다면 json 형태의 key키와 value 값으로 나오기 때문에 곧바로 데이터화 하여 사용할 수 있다.
이렇게 Ollama 모델을 사용하여 질문에서 키워드를 추출할 수 있다. 이 과정은 다양한 질문에 대해 유용하게 활용될 수 있다.