CohereのRerank 3モデルのリリースブログの一番最後に、「Searching Over Tables」という章があり、解説ノートブックのリンクがありました。
リランキングについて理解が深まるかと思い、やってみました。
環境準備
埋め込みにはBedrockのTitan Text Embeddings V2を使用するように改変しています。
Pythonライブラリーをインストールします。
pip install -q pandas numpy fastparquet cohere boto3
Pythonスクリプトを作成します。
まずは、ライブラリーをインポートします。
import json
import boto3
import cohere
import numpy as np
import pandas as pd
Cohereアクセス用のクライアントを生成します。
トライアル用のAPIキーが無料で取得できます。
cohere_key = "<<YOUR_COHERE_API_KEY>>"
co = cohere.Client(cohere_key)
Bedrockアクセス用のクライアントを生成します。
client = boto3.client("bedrock-runtime")
検証に使用するデータセットを取得します。
df = pd.read_parquet("https://huggingface.co/datasets/Cohere/movies/resolve/main/movies.parquet")
以下のようなデータが、4803行含まれています。
| title | overview | genres | producer | cast | |
|---|---|---|---|---|---|
| 0 | Avatar | In the 22nd century, a paraplegic Marine is di... | Action, Adventure, Fantasy, Science Fiction | Ingenious Film Partners, Twentieth Century Fox... | Sam Worthington as Jake Sully, Zoe Saldana as ... |
| 1 | Titanic | 84 years later, a 101-year-old woman named Ros... | Drama, Romance, Thriller | Paramount Pictures, Twentieth Century Fox Film... | Kate Winslet as Rose DeWitt Bukater, Leonardo ... |
| 2 | The Avengers | When an unexpected enemy emerges and threatens... | Science Fiction, Action, Adventure | Paramount Pictures, Marvel Studios | Robert Downey Jr. as Tony Stark / Iron Man, Ch... |
| 3 | Jurassic World | Twenty-two years after the events of Jurassic ... | Action, Adventure, Science Fiction, Thriller | Universal Studios, Amblin Entertainment, Legen... | Chris Pratt as Owen Grady, Bryce Dallas Howard... |
| 4 | Furious 7 | Deckard Shaw seeks revenge against Dominic Tor... | Action | Universal Pictures, Original Film, Fuji Televi... | Vin Diesel as Dominic Toretto, Paul Walker as ... |
| ... | ... | ... | ... | ... | ... |
| 4798 | Midnight Cabaret | A Broadway producer puts on a play with a Devi... | Horror | Lisa Hart Carroll as Dawn, Michael Des Barres ... | |
| 4799 | Growing Up Smith | In 1979, an Indian family moves to America wit... | Comedy, Family, Drama | Roni Akurati as Smith Bhatnagar, Brighton Shar... | |
| 4800 | 8 Days | After sneaking to a party with her friends, 16... | Thriller, Drama | After Eden Pictures | Nicole Smolen as Amber, Kim Baldwin as BB, Ari... |
| 4801 | Running Forever | After being estranged since her mother's death... | Family | New Kingdom Pictures | |
| 4802 | To Be Frank, Sinatra at 100 | None | Documentary | Eyeline Entertainment | Tony Oppedisano as Narrator |
4803 rows × 5 columns
テストデータをベクトル化する
このデータを、Titan Text Embeddings V2でベクトル化します。
一括で埋め込みを行う関数を作成します。
def embed(
texts: list[str],
model_id: str = "amazon.titan-embed-text-v2:0",
):
embeddings = []
for text in texts:
body = {"inputText": text}
response = client.invoke_model(modelId=model_id, body=json.dumps(body))
embeddings.append(json.loads(response["body"].read().decode())["embedding"])
return embeddings
テストデータを、文字列型に変換します。
docs_txt = [str(doc) for doc in df.to_dict("records")]
docs_txt[0]
{'title': 'Avatar', 'overview': 'In the 22nd century, a paraplegic Marine is dispatched to the moon Pandora on a unique mission, but becomes torn between following orders and protecting an alien civilization.', 'genres': 'Action, Adventure, Fantasy, Science Fiction', 'producer': 'Ingenious Film Partners, Twentieth Century Fox Film Corporation, Dune Entertainment, Lightstorm Entertainment', 'cast': "Sam Worthington as Jake Sully, Zoe Saldana as Neytiri, Sigourney Weaver as Dr. Grace Augustine, Stephen Lang as Col. Quaritch, Michelle Rodriguez as Trudy Chacon, Giovanni Ribisi as Selfridge, Joel David Moore as Norm Spellman, CCH Pounder as Moat, Wes Studi as Eytukan, Laz Alonso as Tsu'Tey, Dileep Rao as Dr. Max Patel, Matt Gerald as Lyle Wainfleet, Sean Anthony Moran as Private Fike, Jason Whyte as Cryo Vault Med Tech, Scott Lawrence as Venture Star Crew Chief, Kelly Kilgour as Lock Up Trooper, James Patrick Pitt as Shuttle Pilot, Sean Patrick Murphy as Shuttle Co-Pilot, Peter Dillon as Shuttle Crew Chief, Kevin Dorman as Tractor Operator / Troupe, Kelson Henderson as Dragon Gunship Pilot, David Van Horn as Dragon Gunship Gunner, Jacob Tomuri as Dragon Gunship Navigator, Michael Blain-Rozgay as Suit #1, Jon Curry as Suit #2, Luke Hawker as Ambient Room Tech, Woody Schultz as Ambient Room Tech / Troupe, Peter Mensah as Horse Clan Leader, Sonia Yee as Link Room Tech, Jahnel Curfman as Basketball Avatar / Troupe, Ilram Choi as Basketball Avatar, Kyla Warren as Na'vi Child, Lisa Roumain as Troupe, Debra Wilson as Troupe, Chris Mala as Troupe, Taylor Kibby as Troupe, Jodie Landau as Troupe, Julie Lamm as Troupe, Cullen B. Madden as Troupe, Joseph Brady Madden as Troupe, Frankie Torres as Troupe, Austin Wilson as Troupe, Sara Wilson as Troupe, Tamica Washington-Miller as Troupe, Lucy Briant as Op Center Staff, Nathan Meister as Op Center Staff, Gerry Blair as Op Center Staff, Matthew Chamberlain as Op Center Staff, Paul Yates as Op Center Staff, Wray Wilson as Op Center Duty Officer, James Gaylyn as Op Center Staff, Melvin Leno Clark III as Dancer, Carvon Futrell as Dancer, Brandon Jelkes as Dancer, Micah Moch as Dancer, Hanniyah Muhammad as Dancer, Christopher Nolen as Dancer, Christa Oliver as Dancer, April Marie Thomas as Dancer, Bravita A. Threatt as Dancer, Colin Bleasdale as Mining Chief (uncredited), Mike Bodnar as Veteran Miner (uncredited), Matt Clayton as Richard (uncredited), Nicole Dionne as Nav'i (uncredited), Jamie Harrison as Trooper (uncredited), Allan Henry as Trooper (uncredited), Anthony Ingruber as Ground Technician (uncredited), Ashley Jeffery as Flight Crew Mechanic (uncredited), Dean Knowsley as Samson Pilot, Joseph Mika-Hunt as Trooper (uncredited), Terry Notary as Banshee (uncredited), Kai Pantano as Soldier (uncredited), Logan Pithyou as Blast Technician (uncredited), Stuart Pollock as Vindum Raah (uncredited), Raja as Hero (uncredited), Gareth Ruck as Ops Centreworker (uncredited), Rhian Sheehan as Engineer (uncredited), T. J. Storm as Col. Quaritch's Mech Suit (uncredited), Jodie Taylor as Female Marine (uncredited), Alicia Vela-Bailey as Ikran Clan Leader (uncredited), Richard Whiteside as Geologist (uncredited), Nikie Zambo as Na'vi (uncredited), Julene Renee as Ambient Room Tech / Troupe"}
埋め込みを行います。
doc_emb = np.asarray(embed(texts=docs_txt))
ベクトルデータを使って検索する
「Action movies with Christian Bale」という検索クエリで検索します。
検索クエリをベクトル化します。
query = "Action movies with Christian Bale"
query_emb = np.asarray(embed(texts=[query]))
検索クエリのベクトルと、すべてのドキュメントのベクトルのコサイン類似度/ドット積を計算します。
dot_scores = np.matmul(query_emb, doc_emb.transpose())[0]
上位100件を取得します。
top_100_hits = np.argsort(-dot_scores).tolist()[0:100]
emb_results = df.iloc[top_100_hits].copy()
emb_results.insert(loc=0, column='rank', value=(np.arange(len(emb_results))+1))
emb_results
上位100は、このようになりました。
| rank | title | overview | genres | producer | cast | |
|---|---|---|---|---|---|---|
| 867 | 1 | Babel | Tragedy strikes a married couple on vacation i... | Drama | Paramount Vantage, Media Rights Capital, Anony... | Brad Pitt as Richard, Cate Blanchett as Susan,... |
| 1615 | 2 | Cellular | A young man receives an emergency phone call o... | Action, Adventure, Crime, Thriller | New Line Cinema | Chris Evans as Ryan Hewitt, Kim Basinger as Je... |
| 885 | 3 | Blade | When Blade's mother was bitten by a vampire du... | Horror, Action | New Line Cinema, Amen Ra Films, Imaginary Forc... | Wesley Snipes as Blade, Stephen Dorff as Deaco... |
| 2361 | 4 | Action Jackson | Vengence drives a tough Detroit cop to stay on... | Action, Adventure, Comedy, Crime, Drama | Silver Pictures, Lorimar Motion Pictures | Carl Weathers as Sgt. Jericho 'Action' Jackson... |
| 4072 | 5 | Malone | Erstwhile C.I.A. assassin Richard Malone hopes... | Action, Thriller | Orion Pictures | Burt Reynolds as Richard Malone, Cliff Roberts... |
| ... | ... | ... | ... | ... | ... | ... |
| 2223 | 96 | Desperado | A gunslinger is embroiled in a war with a loca... | Thriller, Action, Crime | Columbia Pictures, Los Hooligans Productions | Antonio Banderas as El Mariachi, Salma Hayek a... |
| 56 | 97 | Spider-Man 2 | Peter Parker is going through a major identity... | Action, Adventure, Fantasy | Columbia Pictures, Laura Ziskin Productions, M... | Tobey Maguire as Peter Parker / Spider-Man, Ki... |
| 1386 | 98 | The Core | Geophysicist Dr. Josh Keyes discovers that an ... | Action, Thriller, Adventure, Science Fiction | Paramount Pictures, David Foster Productions, ... | Aaron Eckhart as Dr. Josh Keyes, Hilary Swank ... |
| 547 | 99 | White House Down | Capitol Policeman John Cale has just been deni... | Action, Drama, Thriller | Columbia Pictures, Centropolis Entertainment, ... | Channing Tatum as Cale, Jamie Foxx as Presiden... |
| 3711 | 100 | The Pink Panther | When the coach of the France soccer team is ki... | Action, Comedy, Crime, Mystery, Family | Columbia Pictures, International Production Co... | Steve Martin as Inspector Jacques Clouseau, Ke... |
100 rows × 6 columns
リランキングして検索結果を改善させる
上位100件のデータをJSONオブジェクトに変換します。
docs = emb_results.to_dict('records')
Rerank3では、JSONオブジェクトのリランク対象のキーを指定できます。今回は、すべてのキーを対象とします。
rank_fields = list(docs[0].keys())
Rerank3モデルでリランキングします。
パラメーターは以下のように指定します。
- query: 検索クエリ（"Action movies with Christian Bale"）
- documents: 上位100件のデータ（1件ごとにJSONオブジェクトになっている）
- top_n: 結果として取得する件数
- model: CohereのモデルID
- rank_fields: リランク対象のフィールド。今回は、JSONの全部のキーを指定
results = co.rerank(query=query, documents=docs, top_n=10, model='rerank-english-v3.0', rank_fields=rank_fields)
埋め込みで検索した上位100件のデータから、リランク処理で上位10件となったデータを抽出します。
top_ids = [hit.index for hit in results.results]
rerank_results = emb_results.iloc[top_ids].copy()
rerank_results['rank'] = (np.arange(len(rerank_results))+1)
結果はこのとおりです。
| rank | title | overview | genres | producer | cast | |
|---|---|---|---|---|---|---|
| 2911 | 1 | Harsh Times | Jim Davis is an ex-Army Ranger who finds himse... | Crime, Drama, Thriller, Action | The Weinstein Company | Christian Bale as Jim Luther Davis, Freddy Rod... |
| 1818 | 2 | Reign of Fire | In post-apocalyptic England, an American volun... | Adventure, Action, Fantasy | The Zanuck Company, Spyglass Entertainment, Wo... | Christian Bale as Quinn Abercromby, Matthew Mc... |
| 2942 | 3 | Equilibrium | In a dystopian future, a totalitarian regime m... | Action, Science Fiction, Thriller | Dimension Films, Blue Tulip Productions | Christian Bale as John Preston, Taye Diggs as ... |
| 2849 | 4 | Rescue Dawn | A US Fighter pilot's epic struggle of survival... | Adventure, Drama, War | Metro-Goldwyn-Mayer (MGM) | Christian Bale as Dieter Dengler, Steve Zahn a... |
| 2528 | 5 | Out of the Furnace | Two brothers live in the economically-depresse... | Thriller, Drama, Crime | Appian Way, Scott Free Productions, Relativity... | Christian Bale as Russell Baze, Zoe Saldana as... |
| 1062 | 6 | Shaft | New York police detective John Shaft arrests W... | Action, Adventure, Crime, Thriller | Paramount Pictures, Scott Rudin Productions | Samuel L. Jackson as John Shaft, Jeffrey Wrigh... |
| 37 | 7 | Spectre | A cryptic message from Bond’s past sends him o... | Action, Adventure, Crime | Columbia Pictures, Danjaq, B24 | Daniel Craig as James Bond, Christoph Waltz as... |
| 2012 | 8 | American Psycho | A wealthy New York investment banking executiv... | Thriller, Drama, Crime | Lions Gate Films, Edward R. Pressman Film, Mus... | Christian Bale as Patrick Bateman, Willem Dafo... |
| 2805 | 9 | The Machinist | The Machinist is the story of Trevor Reznik, a... | Thriller, Drama | Filmax Entertainment, Castelao Producciones, C... | Christian Bale as Trevor Reznik, Jennifer Jaso... |
| 3427 | 10 | Batman | The Dynamic Duo faces four super-villains who ... | Family, Adventure, Comedy, Science Fiction, Crime | Twentieth Century Fox Film Corporation, DC Com... | Adam West as Batman / Bruce Wayne, Burt Ward a... |
ベクトル検索のみで取得した上位5件と、その後リランキングした上位5件を比較してみます。
まず、ベクトル検索のみの結果ですが、castに「Christian Bale」がいません。なので、全部検索結果としてはイケてないです。
対して、リランキングした結果は、すべてのcastに「Christian Bale」が含まれています。
また、「genres」についても、ベクトル検索のみ一位は「Drama」で、Actionではありませんでした。なぜこの結果が一位だったのかは気になりますね。
対して、リランキングした結果は、一位から三位までは「Action」が含まれていました。ただ、四位と五位は違うジャンルでした。ただ、AdventureやThrillerなどが含まれているので広く見てActionと捉えてもいいかもしれない？ように思いました。
上手に関連性の高いものが上位になっているように見えます。
ベクトル検索は、「意味が近いもの」を取ってくるようで、細かく見ると「近そうなもの」が取れているだけなのjかもしれません？リランキングすると、その「近そうなもの」を「近い順」にしてくれるイメージかな、と理解しました。
以上、リランキングの検証でした。