[PaLM-2] VertexAI와 Neo4j로 영화 추천하기

Toward Data Science에 재미있는 코드가 올라와서 조금 수정해서 따라했다. 그래프DB 도구인 Neo4j의 컴퍼넌트인 APOC(A Package of Component)이 OpenAI API와 GCP Vertex AI API를 지원하면서 해당 플랫폼이 지원하는 LLM을 사용할 수 있게 된 것이다. 일반 데이터 베이스와 연동하면 드러나지 않는 데이터의 관계성이 그래프DB에서는 명시적으로 드러나 있으니 연관된 정보를 좀 더 잘 대답해 줄 수 있지 않을까 하는 아이디어에서 시작했다고. 

 

Integrate LLM workflows with Knowledge Graph using Neo4j and APOC

 

그래프DB를 잘 모르지만, 아래와 같이 데이터 간에 관계성을 모두 기록해 두어서 노드 간의 에지를 따라가면 연관된 정보를 찾을 수 있다. 여기서는 Neo4j가 제공하는 작은 데이터 세트(약 30여개의 영화로 구성)인 movie dataset으로  간단한 추천시스템을 만들어 본다.

 

 

아래는 Neo4j 연결을 하고 Query처리하는 간단한 함수를 구성하고 있다(pip install neo4j로 패키지 사전 설치). 

 

# Define Neo4j connections
from neo4j import GraphDatabase
host = 'bolt://3.216.125.181:7687'
user = 'neo4j'
password = 'strikes-welder-laundry'
driver = GraphDatabase.driver(host,auth=(user, password))

def run_query(query, params={}):
    with driver.session() as session:
        result = session.run(query, params)
        return result.to_df()

 

참고로 Neo4j는 데스크톱용으로 설치해서 써도 되지만, 단순히 (3일 이내) 테스트용이라면 샌드박스도 제공해 준다. 위의 host와 user/password 정보는 https://sandbox.neo4j.com/?usecase=movies 에서 connection details를 보면 확인할 수 있다.

 

 

python으로 아래의 코드를 실행하여 접속이 잘 되는지 확인한다. Neo4j의 Cypher Query Language는 잘 모르기도 하고, 여기서 중요한 내용은 아니므로 아래의 쿼리가 데이터 셋에서 1개의 영화 정보를 가져와서 여러 정보를 보여준다고만 알고 있자.

 

print(run_query("""
    MATCH (m:Movie)
    MATCH (m)-[r:ACTED_IN|DIRECTED]-(t)
    WITH m, type(r) as type, collect(t.name) as names
    WITH m, type+": "+reduce(s="", n IN names | s + n + ", ") as types
    WITH m, collect(types) as contexts
    WITH m, "Movie title: "+ m.title + " year: "+coalesce(m.released,"") +" plot: "+ coalesce(m.tagline,"")+"\n" +
           reduce(s="", c in contexts | s + substring(c, 0, size(c)-2) +"\n") as context
    RETURN context LIMIT 1
""")['context'][0])

 

결과 : matrix 영화 정보가 나온다.

Movie title: The Matrix year: 1999 plot: Welcome to the Real World
ACTED_IN: Emil Eifrem, Hugo Weaving, Laurence Fishburne, Carrie-Anne Moss, Keanu Reeves
DIRECTED: Lana Wachowski, Lilly Wachowski

 

참고한 예제와 달리 (현재까지) 무료로 사용 가능한 VertexAI를 사용할 것이므로 아래의 정보를 기록해 둔다. vertexai_access_token은 GCP의 google cloud shell에서 'gcloud auth print-access-token' 명령을 주면 알 수 있고, project_name은 GCP에서 사용하고 있는 project 이름이다.

 

vertexai_access_token = "your_access_token"
project_name = "your_gcp_project_name"

 

이제 movies dataset에 있는 정보들을 GCP VertexAI 내의 Vector Database에 Embedding 시켜준다.

 

run_query("""
    CALL apoc.periodic.iterate(
      'MATCH (m:Movie) RETURN id(m) AS id',
      'MATCH (m:Movie)
       WHERE id(m) = id
       MATCH (m)-[r:ACTED_IN|DIRECTED]-(t)
       WITH m, type(r) as type, collect(t.name) as names
       WITH m, type+": "+reduce(s="", n IN names | s + n + ", ") as types
       WITH m, collect(types) as contexts
       WITH m, "Movie title: "+ m.title + " year: "+coalesce(m.released,"") +" plot: "+ coalesce(m.tagline,"")+"\n" +
            reduce(s="", c in contexts | s + substring(c, 0, size(c)-2) +"\n") as context
       CALL apoc.ml.vertexai.embedding([context], $apiKey, $project) YIELD embedding
       SET m.embedding = embedding',
      {batchSize:1, retries:3, params: {apiKey: $apiKey, project: $project}})
""", {'apiKey': vertexai_access_token, 'project': project_name})['errorMessages'][0]

 

이제 본격적으로 추천을 위한 함수들을 만들어 준다. apoc이 제공하는 chat함수를 통해 vertex ai의 LLM인 PaLM-2를 사용하게 되는데, 필요한 정보는 access_token과 project, 그리고 챗봇의 대답 맥락(system_prompt)과 정교한 질문(?)이다. 일단 챗봇의 맥락, user prompt를 생성하는 함수는 아래와 같다.

 

system_prompt = """
You are an assistant that helps to generate text to form nice and human understandable answers based.
The latest prompt contains the information, and you need to generate a human readable response based on the given information.
Make the answer sound as a response to the question. Do not mention that you based the result on the given information.
Do not add any additional information that is not explicitly provided in the latest prompt.
I repeat, do not add any information that is not explicitly given.
"""

def generate_user_prompt(question, context):
    return f"""
       The question is {question}
       Answer the question by using the provided information:
       {context}
   """

 

아래는 입베딩이 저장된 벡터DB에서 질문과 연관된 영화 정보들을 가져오는 부분이다. 이 정보들을 질문과 결합하여 LLM에 넘겨지게 된다. 여기서는 연관된 3개를 가져 온다.

 

def retrieve_context(question, k=3):
    data = run_query(
        """
            // retrieve the embedding of the question
            CALL apoc.ml.vertexai.embedding([$question], $apiKey, $project) YIELD embedding
            // match relevant movies
            MATCH (m:Movie)
            WITH m, gds.similarity.cosine(embedding, m.embedding) AS score
            ORDER BY score DESC
            // limit the number of relevant documents
            LIMIT toInteger($k)
            // retrieve graph context
            MATCH (m)--()--(m1:Movie)
            WITH m,m1, count(*) AS count
            ORDER BY count DESC
            WITH m, apoc.text.join(collect(m1.title)[..3], ", ") AS similarMovies
            MATCH (m)-[r:ACTED_IN|DIRECTED]-(t)
            WITH m, similarMovies, type(r) as type, collect(t.name) as names
            WITH m, similarMovies, type+": "+reduce(s="", n IN names | s + n + ", ") as types
            WITH m, similarMovies, collect(types) as contexts
            WITH m, "Movie title: "+ m.title + " year: "+coalesce(m.released,"") +" plot: "+ coalesce(m.tagline,"")+"\n" +
                  reduce(s="", c in contexts | s + substring(c, 0, size(c)-2) +"\n") + "similar movies:" + similarMovies + "\n" as context
            RETURN context
          """,
          {"question": question, "k": k, "apiKey": vertexai_access_token, "project": project_name},
    )
    return data["context"].to_list()

 

마지막으로 위에서 가져온 질문과 영화 정보, 그리고 시스템 프롬프트에 기반하여 대답을 LLM Chat 모델이 생성하는 부분이다.

 

def generate_answer(question):
    # Retrieve context
    context = retrieve_context(question)
    # Print context
    for c in context:
        print(c)
    # Generate answer
    response = run_query(
        """
          CALL apoc.ml.vertexai.chat(
          /*messages*/
          [{author:"user", content:$content}],
          $apiKey, $project,
          {temperature:0},
          /*context*/ $system_prompt,
          /*examples*/ [])
          yield value
        """,
        {
            "system_prompt": system_prompt,
            "content": generate_user_prompt(question, context),
            "apiKey": vertexai_access_token,
            "project": project_name
        },
    )
    return response["value"].values[0]['candidates'][0]['content']

 

return 부분이 좀 복잡하게 되어 있는데, response["value"].values[0] 부분이 아래와 같이 부가정보와 결합되어 있기 때문이다. 실질적 대답은 'content' 부분이므로 해당 부분만 리턴 받으면 된다.

 

{'candidates': [{'author': '1',
   'content': 'I would recommend the movie "One Flew Over the Cuckoo\'s Nest". It is a 1975 American drama film directed by Milos Forman and starring Jack Nicholson, Danny DeVito, and Louise Fletcher. The film is based on Ken Kesey\'s 1962 novel of the same name. It tells the story of Randle Patrick McMurphy (Nicholson), a rebellious patient who arrives at a mental institution and challenges the authority of the head nurse, Mildred Ratched (Fletcher). The film was a critical and commercial success, winning five Academy Awards, including Best Picture, Best Director, Best Actor (Nicholson), Best Actress (Fletcher), and Best Adapted Screenplay.'}],
 'citationMetadata': [{'citations': []}],
 'safetyAttributes': [{'blocked': False,
   'scores': [0.1],
   'categories': ['Toxic']}]}

 

이제 실제 테스트를 진행해 보자.

 

[질문]
generate_answer("Who played in the Matrix?")

[대답]
The Matrix was played by Emil Eifrem, Hugo Weaving, Laurence Fishburne, 
Carrie-Anne Moss, and Keanu Reeves.

[질문]
generate_answer("Recommend a movie with Jack Nicholson?")

[대답]
I would recommend the movie One Flew Over the Cuckoo`s Nest. It is a 1975 
American drama film directed by Milos Forman and starring Jack Nicholson, 
Danny DeVito, and Louise Fletcher. The film is based on Ken Kesey`s 1962 
novel of the same name. It tells the story of Randle Patrick McMurphy (Nicholson), 
a rebellious patient who arrives at a mental institution and challenges the 
authority of the head nurse, Mildred Ratched (Fletcher). 
The film was a critical and commercial success, winning five Academy Awards, 
including Best Picture, Best Director, Best Actor (Nicholson), 
Best Actress (Fletcher), and Best Adapted Screenplay.

[질문]
generate_answer("What are similar movies to As Good as It Gets?")

[대답]
Similar movies to As Good as It Gets include A Few Good Men, Cast Away, 
Twister, Joe Versus the Volcano, What Dreams May Come, Bicentennial Man, 
The Birdcage, and Jerry Maguire.

 

생각보다 훌륭히 대답하는데, LLM 자체가 갖고 있는 정보까지 포함하여 대답하고 있다.