Demo of Rerank

1
!pip install "cohere<5"

DEPRECATION: Configuring installation scheme with distutils config files is deprecated and will no longer work in the near future. If you are using a Homebrew or Linuxbrew Python, please see discussion at https://github.com/Homebrew/homebrew-core/issues/76621
Requirement already satisfied: cohere<5 in /opt/homebrew/lib/python3.9/site-packages (4.45)
Requirement already satisfied: aiohttp<4.0,>=3.0 in /opt/homebrew/lib/python3.9/site-packages (from cohere<5) (3.8.1)
Requirement already satisfied: backoff<3.0,>=2.0 in /opt/homebrew/lib/python3.9/site-packages (from cohere<5) (2.2.1)
Requirement already satisfied: fastavro<2.0,>=1.8 in /opt/homebrew/lib/python3.9/site-packages (from cohere<5) (1.9.3)
Requirement already satisfied: importlib_metadata<7.0,>=6.0 in /opt/homebrew/lib/python3.9/site-packages (from cohere<5) (6.6.0)
Requirement already satisfied: requests<3.0.0,>=2.25.0 in /Users/elliottchoi/Library/Python/3.9/lib/python/site-packages (from cohere<5) (2.28.2)
Requirement already satisfied: urllib3<3,>=1.26 in /Users/elliottchoi/Library/Python/3.9/lib/python/site-packages (from cohere<5) (1.26.14)
Requirement already satisfied: attrs>=17.3.0 in /opt/homebrew/lib/python3.9/site-packages (from aiohttp<4.0,>=3.0->cohere<5) (22.1.0)
Requirement already satisfied: charset-normalizer<3.0,>=2.0 in /opt/homebrew/lib/python3.9/site-packages (from aiohttp<4.0,>=3.0->cohere<5) (2.0.12)
Requirement already satisfied: multidict<7.0,>=4.5 in /opt/homebrew/lib/python3.9/site-packages (from aiohttp<4.0,>=3.0->cohere<5) (6.0.2)
Requirement already satisfied: async-timeout<5.0,>=4.0.0a3 in /opt/homebrew/lib/python3.9/site-packages (from aiohttp<4.0,>=3.0->cohere<5) (4.0.2)
Requirement already satisfied: yarl<2.0,>=1.0 in /opt/homebrew/lib/python3.9/site-packages (from aiohttp<4.0,>=3.0->cohere<5) (1.8.1)
Requirement already satisfied: frozenlist>=1.1.1 in /opt/homebrew/lib/python3.9/site-packages (from aiohttp<4.0,>=3.0->cohere<5) (1.3.1)
Requirement already satisfied: aiosignal>=1.1.2 in /opt/homebrew/lib/python3.9/site-packages (from aiohttp<4.0,>=3.0->cohere<5) (1.2.0)
Requirement already satisfied: zipp>=0.5 in /opt/homebrew/lib/python3.9/site-packages (from importlib_metadata<7.0,>=6.0->cohere<5) (3.15.0)
Requirement already satisfied: idna<4,>=2.5 in /Users/elliottchoi/Library/Python/3.9/lib/python/site-packages (from requests<3.0.0,>=2.25.0->cohere<5) (3.4)
Requirement already satisfied: certifi>=2017.4.17 in /Users/elliottchoi/Library/Python/3.9/lib/python/site-packages (from requests<3.0.0,>=2.25.0->cohere<5) (2022.12.7)
DEPRECATION: Configuring installation scheme with distutils config files is deprecated and will no longer work in the near future. If you are using a Homebrew or Linuxbrew Python, please see discussion at https://github.com/Homebrew/homebrew-core/issues/76621


1
import cohere
2
import requests
3
import numpy as np
4
from time import time
5
from typing import List
6
from pprint import pprint

1
API_KEY = "<insert api="" key="" your="">"
2
co = cohere.Client(API_KEY)
3
MODEL_NAME = "rerank-english-v3.0" # another option is rerank-multilingual-02
4
5
query = "What is the capital of the United States?"
6
docs = [
7
    "Carson City is the capital city of the American state of Nevada. At the 2010 United States Census, Carson City had a population of 55,274.",
8
    "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean that are a political division controlled by the United States. Its capital is Saipan.",
9
    "Charlotte Amalie is the capital and largest city of the United States Virgin Islands. It has about 20,000 people. The city is on the island of Saint Thomas.",
10
    "Washington, D.C. (also known as simply Washington or D.C., and officially as the District of Columbia) is the capital of the United States. It is a federal district. The President of the USA and many major national government offices are in the territory. This makes it the political center of the United States of America.",
11
    "West Virginia is a state in the Appalachian region of the United States. Its capital and largest city is Charleston. It is often abbreviated W. Va. or simply WV.",
12
    "Capital punishment (the death penalty) has existed in the United States since before the United States was a country. As of 2017, capital punishment is legal in 30 of the 50 states. The federal government (including the United States military) also uses capital punishment.",
13
    "North Dakota is a state in the United States. 672,591 people lived in North Dakota in the year 2010. The capital and seat of government is Bismarck.",
14
    "Kentucky is a state in the United States. Its capital is Frankfort. It touches the states of Missouri (by the Mississippi River), Illinois, Indiana, Ohio, West Virginia (by the Ohio River), Tennessee and Virginia. There are many rivers in Kentucky",
15
    "Micronesia, officially the Federated States of Micronesia, is an island nation in the Pacific Ocean, northeast of Papua New Guinea. The country is a sovereign state in free association with the United States. The capital city of Federated States of Micronesia is Palikir.",
16
    "Utah is a state in the west United States. The capital and largest city is Salt Lake City. Utah became a state in the U.S. on January 4, 1896."]

1
results = co.rerank(query=query, model=MODEL_NAME, documents=docs, top_n=3) # Change top_n to change the number of results returned. If top_n is not passed, all results will be returned.
2
for idx, r in enumerate(results):
3
  print(f"Document Rank: {idx + 1}, Document Index: {r.index}")
4
  print(f"Document: {r.document['text']}")
5
  print(f"Relevance Score: {r.relevance_score:.2f}")
6
  print("\n")

Document Rank: 1, Document Index: 3
Document: Washington, D.C. (also known as simply Washington or D.C., and officially as the District of Columbia) is the capital of the United States. It is a federal district. The President of the USA and many major national government offices are in the territory. This makes it the political center of the United States of America.
Relevance Score: 1.00
Document Rank: 2, Document Index: 5
Document: Capital punishment (the death penalty) has existed in the United States since before the United States was a country. As of 2017, capital punishment is legal in 30 of the 50 states. The federal government (including the United States military) also uses capital punishment.
Relevance Score: 0.75
Document Rank: 3, Document Index: 1
Document: The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean that are a political division controlled by the United States. Its capital is Saipan.
Relevance Score: 0.09

1
!pip install -U  rank_bm25

DEPRECATION: Configuring installation scheme with distutils config files is deprecated and will no longer work in the near future. If you are using a Homebrew or Linuxbrew Python, please see discussion at https://github.com/Homebrew/homebrew-core/issues/76621
Collecting rank_bm25
    Downloading rank_bm25-0.2.2-py3-none-any.whl.metadata (3.2 kB)
Requirement already satisfied: numpy in /opt/homebrew/lib/python3.9/site-packages (from rank_bm25) (1.23.5)
Downloading rank_bm25-0.2.2-py3-none-any.whl (8.6 kB)
Installing collected packages: rank_bm25
  DEPRECATION: Configuring installation scheme with distutils config files is deprecated and will no longer work in the near future. If you are using a Homebrew or Linuxbrew Python, please see discussion at https://github.com/Homebrew/homebrew-core/issues/76621
DEPRECATION: Configuring installation scheme with distutils config files is deprecated and will no longer work in the near future. If you are using a Homebrew or Linuxbrew Python, please see discussion at https://github.com/Homebrew/homebrew-core/issues/76621
Successfully installed rank_bm25-0.2.2

1
import json
2
import gzip
3
import os
4
from rank_bm25 import BM25Okapi
5
from sklearn.feature_extraction import _stop_words
6
import string
7
from tqdm.autonotebook import tqdm

/var/folders/ww/ht8qwj2s7s799qnktblg6qhm0000gp/T/ipykernel_31832/1066443236.py:7: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
    from tqdm.autonotebook import tqdm

1
!wget http://sbert.net/datasets/simplewiki-2020-11-01.jsonl.gz

--2024-04-08 14:28:00--  http://sbert.net/datasets/simplewiki-2020-11-01.jsonl.gz
Resolving sbert.net (sbert.net)... 172.64.80.1, 2606:4700:130:436c:6f75:6466:6c61:7265
Connecting to sbert.net (sbert.net)|172.64.80.1|:80... connected.
HTTP request sent, awaiting response... 301 Moved Permanently
Location: https://sbert.net/datasets/simplewiki-2020-11-01.jsonl.gz [following]
--2024-04-08 14:28:01--  https://sbert.net/datasets/simplewiki-2020-11-01.jsonl.gz
Connecting to sbert.net (sbert.net)|172.64.80.1|:443... connected.
HTTP request sent, awaiting response... 301 Moved Permanently
Location: https://public.ukp.informatik.tu-darmstadt.de/reimers/sentence-transformers/datasets/simplewiki-2020-11-01.jsonl.gz [following]
--2024-04-08 14:28:01--  https://public.ukp.informatik.tu-darmstadt.de/reimers/sentence-transformers/datasets/simplewiki-2020-11-01.jsonl.gz
Resolving public.ukp.informatik.tu-darmstadt.de (public.ukp.informatik.tu-darmstadt.de)... 130.83.167.186
Connecting to public.ukp.informatik.tu-darmstadt.de (public.ukp.informatik.tu-darmstadt.de)|130.83.167.186|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 50223724 (48M) [application/octet-stream]
Saving to: ‘simplewiki-2020-11-01.jsonl.gz’
simplewiki-2020-11- 100%[===================>]  47.90M  5.78MB/s    in 8.9s
2024-04-08 14:28:11 (5.37 MB/s) - ‘simplewiki-2020-11-01.jsonl.gz’ saved [50223724/50223724]

1
wikipedia_filepath = 'simplewiki-2020-11-01.jsonl.gz'
2
3
passages = []
4
with gzip.open(wikipedia_filepath, 'rt', encoding='utf8') as fIn:
5
    for line in fIn:
6
        data = json.loads(line.strip())
7
        passages.extend(data['paragraphs'])
8
9
print("Passages:", len(passages))

Passages: 509663

1
print(passages[0], passages[1])

1
def bm25_tokenizer(text):
2
    tokenized_doc = []
3
    for token in text.lower().split():
4
        token = token.strip(string.punctuation)
5
6
        if len(token) > 0 and token not in _stop_words.ENGLISH_STOP_WORDS:
7
            tokenized_doc.append(token)
8
    return tokenized_doc
9
10
11
tokenized_corpus = []
12
for passage in tqdm(passages):
13
    tokenized_corpus.append(bm25_tokenizer(passage))
14
15
bm25 = BM25Okapi(tokenized_corpus)

1
def search(query, top_k=3, num_candidates=100):
2
    print("Input question:", query)
3
4
    ##### BM25 search (lexical search) #####
5
    bm25_scores = bm25.get_scores(bm25_tokenizer(query))
6
    top_n = np.argpartition(bm25_scores, -num_candidates)[-num_candidates:]
7
    bm25_hits = [{'corpus_id': idx, 'score': bm25_scores[idx]} for idx in top_n]
8
    bm25_hits = sorted(bm25_hits, key=lambda x: x['score'], reverse=True)
9
10
    print(f"Top-3 lexical search (BM25) hits")
11
    for hit in bm25_hits[0:top_k]:
12
        print("\t{:.3f}\t{}".format(hit['score'], passages[hit['corpus_id']].replace("\n", " ")))
13
14
15
    #Add re-ranking
16
    docs = [passages[hit['corpus_id']] for hit in bm25_hits]
17
18
    print(f"\nTop-3 hits by rank-API ({len(bm25_hits)} BM25 hits re-ranked)")
19
    results = co.rerank(query=query, model=MODEL_NAME, documents=docs, top_n=top_k)
20
    for hit in results:
21
        print("\t{:.3f}\t{}".format(hit.relevance_score, hit.document["text"].replace("\n", " ")))

1
search(query = "What is the capital of the United States?")

Input question: What is the capital of the United States?
Top-3 lexical search (BM25) hits
    16.264	Capital punishment (the death penalty) has existed in the United States since before the United States was a country. As of 2017, capital punishment is legal in 30 of the 50 states. The federal government (including the United States military) also uses capital punishment.
    15.124	In 1783, it was the capital of the United States for a few months.
    14.476	New York was the capital of the United States under the Articles of Confederation from 1785 to 1788. When the US Constitution was made, it stayed as the capital from 1789 until 1790. In 1789, the first President of the United States, George Washington, was inaugurated; the first United States Congress and the Supreme Court of the United States each met for the first time, and the United States Bill of Rights was written, all at Federal Hall on Wall Street. By 1790, New York grew bigger than Philadelphia, so it become the biggest city in the United States. By the end of 1790, because of the Residence Act, Philadelphia became the new capital.
Top-3 hits by rank-API (100 BM25 hits re-ranked)
    0.999	Washington, D.C. (also known as simply Washington or D.C., and officially as the District of Columbia) is the capital of the United States. It is a federal district. The President of the USA and many major national government offices are in the territory. This makes it the political center of the United States of America.
    0.994	New York was the capital of the United States under the Articles of Confederation from 1785 to 1788. When the US Constitution was made, it stayed as the capital from 1789 until 1790. In 1789, the first President of the United States, George Washington, was inaugurated; the first United States Congress and the Supreme Court of the United States each met for the first time, and the United States Bill of Rights was written, all at Federal Hall on Wall Street. By 1790, New York grew bigger than Philadelphia, so it become the biggest city in the United States. By the end of 1790, because of the Residence Act, Philadelphia became the new capital.
    0.993	As the national capital of the United States, Washington, D.C. has numerous media outlets in various mediums. Some of these media are known throughout the United States, including "The Washington Post" and various broadcasting networks headquartered in D.C.

1
search(query = "Number countries Europe")

Input question: Number countries Europe
Top-3 lexical search (BM25) hits
    16.963	ECoHR' has a number of judges. The number of judges is seven normally but at the case of dealing a great issue, the number will be 21 and the judges are equally from member countries of the Council of Europe. At present, there are forty seven member countries of the Council of Europe. Each country may have one judge in the ECoHR. But, judges work independently for the ECoHR, and not for their country.
    14.560	Most countries in Europe, and a few countries in Asia, have made some or all synthetic cannabinoids illegal.
    14.165	Many of these countries were members of the Western European Union. Many, such as Norway, are also in Northern Europe or in Central Europe or Southern Europe.
Top-3 hits by rank-API (100 BM25 hits re-ranked)
    0.997	There are at least 43 countries in Europe (the European identities of 5 transcontinental countries:Cyprus, Georgia, Kazakhstan, Russia and Turkey are disputed). Most of these countries are members of the European Union.
    0.987	Within these regions, there are up to 48 independent European countries (with the identities of 5 transcontinental countries being disputed). The largest is the Russian Federation, which covers 39% of Europe.
    0.981	Europe, the planet's 6th largest continent, includes 47 countries and assorted dependencies, islands and territories.

1
search(query = "Elon Musk year birth")

Input question: Elon Musk year birth
Top-3 lexical search (BM25) hits
    22.568	Tesla, Inc. is a company based in Palo Alto, California which makes electric cars. It was started in 2003 by Martin Eberhard, Dylan Stott, and Elon Musk (who also co-founded PayPal and SpaceX and is the CEO of SpaceX). Eberhard no longer works there. Today, Elon Musk is the Chief Executive Officer (CEO). It started selling its first car, the Roadster in 2008.
    20.492	Elon Musk complained via Twitter about Los Angeles traffic and the same day, December 17, 2016, founded the company. It built a short test tunnel in Los Angeles.
    20.448	At the end of 2016, Musk founded The Boring Company which focuses on tunnelling and infrastructure. He mentioned Los Angeles traffic as the reason for starting this company. In March 2017 Elon Musk announced he has started another company which aims to merge human brains and computers, it is called Neuralink.
Top-3 hits by rank-API (100 BM25 hits re-ranked)
    0.994	Elon Reeve Musk (born June 28, 1971) is a businessman and philanthropist. He was born in South Africa. He moved to Canada and later became an American citizen. Musk is the current CEO & Chief Product Architect of Tesla Motors, a company that makes electric vehicles. He is also the CEO of Solar City, a company that makes solar panels, and the CEO & CTO of SpaceX, an aerospace company. In August 2020, Bloomberg ranked Musk third among the richest people on the planet with net worth to be $115.4 billion.
    0.602	Elon Musk and his brother started Zip2, a software company, in 1995. In 1999 he sold it and became a millionaire. He then started X.com, which merged with the company to make PayPal. X.com was then renamed to PayPal, and he focused on growing that part of the company. He then started SpaceX and became the CEO of Tesla.
    0.474	In early 2002, Musk was seeking workers for his new space company, soon to be named SpaceX. Musk found a rocket engineer Tom Mueller (later SpaceX's CTO of Propulsion). He agreed to work for Musk. That was how SpaceX was born. The first headquarters of SpaceX was in a warehouse in El Segundo, California. The company has grown rapidly since it was founded in 2002, growing from 160 workers in November 2005 to 1,100 in 2010, 3,800 workers and contractors by October 2013, nearly 5,000 by late 2015, and about 6,000 in April 2017.

1
search(query = "Which US president was killed?")

Input question: Which US president was killed?
Top-3 lexical search (BM25) hits
    11.966	He came into office when the previous president, Cyprien Ntaryamira, was killed in a plane crash. It was an assassination in which the Rwandan president Juvénal Habyarimana was also killed. Ntibantunganya left office when he was deposed by Pierre Buyoya in a military coup of 1996.
    11.697	Burr killed Alexander Hamilton in a duel in 1804, when Burr was still Vice President.
    11.482	After President James A. Garfield died, vice-president Chester Arthur replaced him. The man who killed him expected the new President to pardon him. This did not happen.
Top-3 hits by rank-API (100 BM25 hits re-ranked)
    0.984	James Abram Garfield (November 19, 1831 - September 19, 1881) was the 20th (1881) President of the United States and the 2nd President to be assassinated (killed while in office). President Garfield was in office from March to September of 1881. He was in office for a total of six months and fifteen days. For almost half that time he was bedridden as a result of an attempt to kill him. He was shot on July 2 and finally died in September the same year he got into office.
    0.976	President William McKinley was killed by anarchist Leon Czolgosz because Czolgosz believed president McKinley was against good working people, he considered McKinley responsible for falsifying the reasons for the war, and approving and waging an illegal, devastating Philippines war.
    0.916	On the night that President Abraham Lincoln was killed, someone also tried to kill Seward. For the rest of his life, Seward had scars on his face from the attack. Later, the man who attacked him was caught and put to death.

1
search(query="When is Chinese New Year")

Input question: When is Chinese New Year
Top-3 lexical search (BM25) hits
    18.606	Today in China the Gregorian calendar is used for most activities. At the same time, the Chinese calendar is still used for traditional Chinese holidays like Chinese New Year or Lunar New Year.
    18.151	Before that, the holiday was usually just called the "NewYear". Because the traditional Chinese calendar is mostly based on the changes in the moon, the Chinese New Year is also known in English as the "Lunar New Year" or "Chinese Lunar New Year". This name comes from "Luna", an old Latin name for the moon. The Indonesian name for the holiday is Imlek, which comes from the Hokkien word for the old Chinese calendar and is therefore also like saying "Lunar New Year".
    18.011	Spring Festival is the Chinese New Year.
Top-3 hits by rank-API (100 BM25 hits re-ranked)
    0.999	Chinese New Year, known in China as the SpringFestival and in Singapore as the LunarNewYear, is a holiday on and around the new moon on the first day of the year in the traditional Chinese calendar. This calendar is based on the changes in the moon and is only sometimes changed to fit the seasons of the year based on how the Earth moves around the sun. Because of this, Chinese New Year is never on January1. It moves around between January21 and February20.
    0.997	Chinese New Year always starts on a new moon, when the Moon is between the Earth and Sun and it looks all dark in the night sky. Because new moons happen about every 29.53 days but the year set by Pope GregoryXIII is 365.2425 days long, the Chinese holiday moves to different days each year. The Chinese calendar adds a 13th month every so often to keep the seasons in the right place, so the first day of the new year always happens between January21 and February20 on the 2nd or 3rd new moon after the 1st day of winter. The chart on the right gives the day of each Chinese New Year from 1996 to 2031.
    0.996	Chinese New Year lasts fifteen days, including one week as a national holiday. It starts with the first day of the Chinese lunar year and ends with the full moon fifteen days later. It is always in the middle of winter, but is called the Spring Festival in Chinese because Chinese seasons are a little different from English ones. On the first day of the Chinese New Year, people call on friends and relatives. Because most people watch the special performances on CCTV all the night on New Year's Eve and don't go to bed until 12:00 AM, they usually get up later in the next day. The fifth day of the Chinese New Year is the day to welcome the god of Wealth (Chinese:财神爷), many people make and eat dumplings (Chinese:饺子. Pinyin: Jaozi). They believe that dumplings can hold the god of Wealth and bring luck. The last day of the Chinese New Year is the Lantern Festival. On this day, the moon becomes the full moon. People go out and watch the lantern festivals everywhere. After that, they eat sweet dumpling (Chinese:汤圆,元宵), a kind of dumpling which is round and looks like the full moon.

1
search(query="How many people live in Paris")

Input question: How many people live in Paris
Top-3 lexical search (BM25) hits
    16.277	Live à Paris (English: "Live in Paris") is a live album by Canadian singer Céline Dion.
    15.173	Île-de-France is a region of France. The capital city is Paris. It is also the capital city of France. In 2013 about 12 million people lived in the region. About 2.1 million people live in the city of Paris.
    14.666	Gennevilliers is a town in France near Paris. It is in the region Île-de-France and the department of Hauts-de-Seine. About 41,000 people live there.
Top-3 hits by rank-API (100 BM25 hits re-ranked)
    0.999	Paris (nicknamed the ""City of light"") is the capital city of France, and the largest city in France. The area is , and around 2.15 million people live there. If suburbs are counted, the population of the Paris area rises to 12 million people.
    0.987	Île-de-France is a region of France. The capital city is Paris. It is also the capital city of France. In 2013 about 12 million people lived in the region. About 2.1 million people live in the city of Paris.
    0.602	Essonne is a department to the south of Paris in the Île-de-France region. Its prefecture is Évry. About 1,172,000 people live there (2006 estimation).

1
search(query="Who is the director of The Matrix?")

Input question: Who is the director of The Matrix?
Top-3 lexical search (BM25) hits
    16.253	An inverse matrix is a matrix that, when multiplied by another matrix, equals the identity matrix. For example:
    16.072	is an identity matrix. There is exactly one identity matrix for each square dimension set. An identity matrix is special because when multiplying any matrix by the identity matrix, the result is always the original matrix with no change.
    15.353	First, the system needs to be turned into an augmented matrix. In an augmented matrix, each linear equation becomes a row. On one side of the augmented matrix, the coefficients of each term in the linear equation become numbers in the matrix. On the other side of the augmented matrix are the constant terms each linear equation is equal to. For this system, the augmented matrix is:
Top-3 hits by rank-API (100 BM25 hits re-ranked)
    0.995	The Matrix is a science fiction action movie that was made in 1999. It was written and directed by the Wachowski Brothers. The main actors in the movie are Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving. "The Matrix" was followed by two sequels: "The Matrix Reloaded" and "The Matrix Revolutions".
    0.992	Helmut Bakaitis (born 26 September 1944) is a German-born Australian director, actor and screenwriter. He is known for his role as The Architect in "The Matrix" movie series. Bakaitis was born in Lauban, Lower Silesia, Germany (now Lubań, Poland). Bakaitis started teaching directing at Australian Academy of Dramatic Art (AADA).
    0.804	The Matrix Revolutions is a 2003 movie that was written and directed by the Wachowski brothers. It is the sequel to "The Matrix Reloaded".