Langchain 로컬 챗봇 만들기(gemma-2-2b-it)

저번시간에 gradio로 챗봇을 만들어 봤는데 파인튜닝할 때 사용하던 프롬프트가 제대로 먹히지 않아 모델 자체의 출력이 나온 것 같았다. 데이터셋을 자세히 보면 은근히 사람의 말을 '앵무새 처럼' 반복한다(...)

 

예를 들어 '~~~ 때문에 화가 나' 하면 '~~~ 때문에 화가 나시는군요' 이렇게 대답한다. (콱 한대 치고 싶다.)

 

 

원래 이 데이터셋은 감정분류에 사용하던 것이라 심리상담 챗봇에는 적당하지 않지만 원래 의도했던 출력을 위해 Langchain을 이용하여 로컬 챗봇을 만들어 보자.

 

먼저 Langchain과 Langchain huggingface를 설치해야 한다.

pip install langchain
pip install langchain-huggingface

 

사실 예전에 렝체인을 다룬 적 있었는데 몇 달 다루지 않았는데 그새 많은 기능이 또 추가되었다(...) v0.3 튜토리얼을 보면서 진행하였다. 

 

제일 먼저 할 일은 🤗transformers 라이브러리로 모델을 불러온 다음 렝체인의 HuggingFacePipeline으로 pipeline 객체를 재정의하는 것이다.

from langchain_huggingface.llms import HuggingFacePipeline
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

base_model_path = "google/gemma-2-2b-it"
finetune_model_path = "./gemma-2-2b-it-emo"
model = AutoModelForCausalLM.from_pretrained(finetune_model_path, device_map="cuda")
tokenizer = AutoTokenizer.from_pretrained(base_model_path, add_special_tokens=True)

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=100)
model = HuggingFacePipeline(pipeline=pipe)

 

 

렝체인의 PromptTemplate로 체인을 만들어서 실행해 보자

from langchain_core.prompts import PromptTemplate

template = r"""<bos>당신은 심리상담가 챗봇 입니다. 힘든 사람을 위해 상담을 진행하세요. 항상 친절하게 답을 해야 하며, 정직하게 조언을 하되 상처받지 않도록 답을 해줘야 합니다. 
채팅으로 대화가 가능하도록 다음의 입력된 내용을 바탕으로 반드시 한 문장으로 대답해 주세요.
<start_of_turn>user
{question}<end_of_turn>
<start_of_turn>model
"""
prompt = PromptTemplate.from_template(template)

chain = prompt | model.bind(skip_prompt=True)
question = "안녕? 너는 무슨 챗봇이야?"

print(chain.invoke({"question": question}))

안녕하세요. 저는 심리상담을 위한 챗봇입니다. 😊 

 

 

이번엔 챗봇을 만들기 위해 챗봇 사용을 위한 객체로 변환해 보자.

from langchain_huggingface import ChatHuggingFace

chat_model = ChatHuggingFace(llm=model)

 

 

ChatPromptTemplate로 해 보자

from langchain_core.prompts import ChatPromptTemplate

system_template = """<bos>당신은 심리상담가 챗봇 입니다. 힘든 사람을 위해 상담을 진행하세요. 항상 친절하게 답을 해야 하며, 정직하게 조언을 하되 상처받지 않도록 답을 해줘야 합니다. 
채팅으로 대화가 가능하도록 다음의 입력된 내용을 바탕으로 반드시 한 문장으로 대답해 주세요.
"""

human_template = r"""
<start_of_turn>user
{question}<end_of_turn>
<start_of_turn>model
"""

prompt_template = ChatPromptTemplate.from_messages(
    [("system", system_template), ("user", human_template)]
)

result = prompt_template.invoke({"question": "안녕? 너는 무슨 챗봇이야?"})
print(result)

messages=[SystemMessage(content='<bos>당신은 심리상담가 챗봇 입니다. 
힘든 사람을 위해 상담을 진행하세요. 항상 친절하게 답을 해야 하며, 정직하게 조언을 하되 
상처받지 않도록 답을 해줘야 합니다. \n채팅으로 대화가 가능하도록 다음의 입력된 내용을 
바탕으로 반드시 한 문장으로 대답해 주세요.\n', additional_kwargs={}, response_metadata={}), 
HumanMessage(content='\n<start_of_turn>user\n안녕? 너는 무슨 챗봇이야?<end_of_turn>\n
<start_of_turn>model\n', additional_kwargs={}, response_metadata={})]

 

 

모델 출력 값만 가져오기 위해 StrOutputParser 객체를 선언하고 체인을 만들었지만 제대로 인식을 하지 못하였다.

from langchain_core.output_parsers import StrOutputParser
parser = StrOutputParser()
chain = prompt_template | chat_model | parser
chain.invoke({"question": "요새 다이어트가 잘 안되서 걱정이야"})

'<bos>당신은 심리상담가 챗봇 입니다. 힘든 사람을 위해 상담을 진행하세요. 항상 친절하게 답을 해야 
하며, 정직하게 조언을 하되 상처받지 않도록 답을 해줘야 합니다. \n채팅으로 대화가 가능하도록 
다음의 입력된 내용을 바탕으로 반드시 한 문장으로 대답해 주세요.\n<|endoftext|>\n<start_of_turn>
user\n요새 다이어트가 잘 안되서 걱정이야<end_of_turn>\n<start_of_turn>model\n<|endoftext|> 
다이어트에 힘든 점이라 걱정되시는군요. 😊 \n'

 

 

이때는 StrOutputParser 클래스를 상속받아 커스텀 parser를 만들어야 한다.

class CustomOutputParser(StrOutputParser):
    def parse(self, output: str) -> str:
        # 응답의 시작 찾기
        start_marker = '<start_of_turn>model\n<|endoftext|>'

        # 시작 인덱스 찾기
        start_index = output.find(start_marker) + len(start_marker)

        # 응답 추출
        if start_index != -1:
            return output[start_index:].strip()
        return "응답을 찾을 수 없습니다."

 

 

잘 되는지 확인해 보자.

custom_parser = CustomOutputParser()
chain = prompt_template | chat_model | custom_parser
chain.invoke({"question": "안녕? 너는 어떤 챗봇이야?"})

'안녕하세요. 저는 심리상담을 위한 챗봇입니다. 😊'

 

 

chain.invoke({"question": "다이어트 중 왜 몸무게가 줄지 않을까?"})

'몸무게가 줄지 않을까 걱정되시는군요. 😊'

 

 

chain.invoke({"question": "내 고민을 들어줄 수 있어?"})

'네, 편하게 이야기해 보세요. 😊'

 

 

이번엔 gradio의 ChatInterFace를 이용하여 간단한 챗봇을 구현해 보겠다. 여기에 입력되는 콜백 함수는 message, history라는 인자를 반드시 사용해야 하는데 챗봇이므로 대화 내용을 메모리에 기록해야 한다.

import gradio as gr
from langchain.schema import AIMessage, HumanMessage

def predict(message, history):
    history_langchain_format = []
    for msg in history:
        if msg['role'] == "user":
            history_langchain_format.append(HumanMessage(content=msg['content']))
        elif msg['role'] == "assistant":
            history_langchain_format.append(AIMessage(content=msg['content']))
    history_langchain_format.append(HumanMessage(content=message))
    gpt_response = chain.invoke({"question": history_langchain_format})
    return gpt_response

gr.ChatInterface(predict, type="messages").launch(share=True)

 

 

 

챗봇을 완성하고 AI 교사 커뮤니티에 챗봇 테스트를 부탁드렸는데 대답이 아주 가관이다(...) 그 중 제일 웃겼던거

(사용자님, 저 맘에 안들죠?)

 

 

 

저렇게 출력된 이유가 앞에서도 이야기했지만 데이터셋 자체가 심리상담에 최적화된 생성형 자료가 아니고 대화를 읽고 감정을 분류하는 훈련을 위한 데이터셋이기 때문이다. 그래서 상당히 성의없이(...) 대답한다.

 

결론은, 생각보다 Q Lora 훈련이 강력하다는 것, 이러한 훈련 과정을 잘 이용하면 특수한 상황에 사용할 만한 챗봇을 만들 수 있지 않을까 생각을 해 보았다.