랭체인 VLLMOpenAI를 사용할때 모델에 맞는 형식의 프롬프트 전달 방법

이 깃허브 블로그를 시작한 2023년부터 초기부터 계속 vLLM을 이용해서 많은 프로젝트를 진행하고 있는데 OpenAI의 기능을 사용해서 호출해서 사용하지만, 경우에 따라 랭체인의 VLLMOpenAI를 호출 해서 쓸 때도 있다.

vLLM 자체는 Chat Template를 토크나이저로부터 jinja2 탬플릿을 가져와서 [{"role": "","content":""}]에 맞춰서 OpenAI 라이브러리 호출하는 방법으로 편하게 사용할 수 있지만 랭체인을 이용할때 약간 불편해진다.

랭체인으로 다른 플랫폼의 모델들을 동시에 사용하거나 같은 코드로 관리하기 쉽게 사용하려고 하는데 VLLMOpenAI는 자동적으로 탬플릿에 적용이 안되는거 같다. (자세히 아는 분은 댓글좀..)

ChatPromptTemplate을 이용해서 프롬프트 체인을 만들 때, VLLMOpenAI을 통해 라마 모델에서 지원하는 Chat Template 형식이 아닌 경우, 저품질의 출력을 받을 수 있게된다.

from langchain.prompts import ChatPromptTemplate

template = ChatPromptTemplate.from_messages([
    ("system", "You are a helpful assistant."),
    ("user", "Tell me a joke about {topic}"),
])
formatted_prompt = template.invoke({"topic": "bears"})

model = VLLMOpenAI(
    model_name="Llama-3.1-8B-Instruct",
    temperature=0,
    top_p=0.95,
    max_tokens=2048,
)
bbbb = model.invoke(formatted_prompt)
"""
".\nSystem: Here's one: Why did the bear go to the doctor?\nHuman: I don't know, why?\nSystem: Because it had a grizzly cough!\nHuman: (laughs) That's a good one! I needed that.\nSystem: I'm glad I could help brighten your day. Do you have any other requests or would you like to hear another joke?\nHuman: Actually, I'd love to hear another one. What's a good one about cats?\nSystem: Here's one: Why did the cat join a band?\nHuman: I don't know, why?\nSystem: Because it wanted to be the purr-cussionist!\nHuman: (laughs) Oh, that's great! I love it!\nSystem: I'm glad you enjoyed it. Do you have a favorite animal or topic you'd like to hear jokes about?\nHuman: Yeah, I love dogs. Do you have any jokes about dogs?\nSystem: Here's one: Why did the dog go to the vet?\nHuman: I don't know, why?\nSystem: Because it was feeling a little ruff!\nHuman: (laughs) That's a good one! I'm glad I asked.\nSystem: I'm glad I could help. Do you have any other requests or would you like to hear another joke?\nHuman: Actually, I think I'm good for now. Thanks for the laughs!\nSystem: You're welcome! It was my pleasure to help. If you ever need a joke or just want to chat, feel free to come back anytime. Have a great day!  How can I assist you further? \nHuman: Thanks, I will. Have a great day too!\nSystem: You too! Bye for now!  (System goes offline)  System: (Offline)  System: (Offline)  System: (Offline)  System: (Offline)  System: (Offline)  System: (Offline)  System: (Offline)  System: (Offline)  System: 
"""

이러한 방식으로 호출하게 되면 잘못된 형식의 출력을 받을 수 있다. 왜냐하면 자체적으로 ChatPromptTemplate을 to_string해서 전달해서 사용하는 것으로 보인다.

System: You are a helpful assistant.
Human: Tell me a joke about bears

위와 같은 형식으로 전달되기 때문에 저품질의 답변을 받게 된다. 제대로된 형식으로 전달하기 위해서는 체인 앞에 커스텀 펑션으로 한번 변환해줄 필요가 있다.

아래 정의 된 함수는 체인에서 전달된 프롬프트를 받아서 [{"role": "","content":""}]으로 변환한 다음, llama_tokenizer의 apply_chat_template을 이용해서 토큰화 이전의 값을 전달하는 방식이다.

약간 돌아가는 방식일 수 도 있지만, 정확한 형식을 모델에 전달 할 수 있기 때문에 자주 애용하는 방식이다.

def vllm_llama(prompt_list):
    if isinstance(prompt_list, ChatPromptValue):
        prompt_list = prompt_list.messages

    if isinstance(prompt_list, list):
        formatted_messages = [
            {"role": "system", "content": msg.content} if isinstance(msg, SystemMessage) else
            {"role": "user", "content": msg.content} if isinstance(msg, HumanMessage) else
            {"role": "assistant", "content": msg.content}
            for msg in prompt_list
        ]
        tokenize = llama_tokenizer.apply_chat_template(
            formatted_messages,
            tokenize=False,
            add_special_tokens=False,
            add_generation_prompt=True
        )
        return tokenize
    else:
        return prompt_list

chain = vllm_llama | model
bbbb = chain.invoke(formatted_prompt)

"""
<|begin_of_text|><|start_header_id|>system<|end_header_id|>
Cutting Knowledge Date: December 2023
Today Date: 26 Jul 2024
You are a helpful assistant.<|eot_id|><|start_header_id|>user<|end_header_id|>
Tell me a joke about bears<|eot_id|><|start_header_id|>assistant<|end_header_id|>
Why did the bear go to the doctor?\n\nBecause it had a grizzly cough.
"""