Embeddings d'articles Wikipedia pour la recherche

Dans ce tutoriel, nous explorerons le processus de préparation d'un dataset d'articles Wikipedia pour la recherche, en nous concentrant spécifiquement sur les Jeux olympiques d'hiver 2022. Cela implique plusieurs étapes clés : collecter les articles, les découper en sections gérables, créer des embeddings pour ces sections, et stocker les résultats pour une récupération efficace.

Prérequis

Avant de commencer, assurez-vous d'avoir les bibliothèques nécessaires installées. Vous pouvez installer les bibliothèques manquantes avec pip :

pip install openai mwclient mwparserfromhell pandas tiktoken

Définissez votre clé API OpenAI comme variable d'environnement :

export OPENAI_API_KEY="votre-cle-api"

Collecte des articles

Nous commençons par télécharger les articles Wikipedia liés aux Jeux olympiques d'hiver 2022. En utilisant la bibliothèque mwclient, nous pouvons accéder à l'API de Wikipedia pour récupérer les articles d'une catégorie spécifique.

import mwclient
 
CATEGORY_TITLE = "Category:2022 Winter Olympics"
WIKI_SITE = "en.wikipedia.org"
 
site = mwclient.Site(WIKI_SITE)
category_page = site.pages[CATEGORY_TITLE]
titles = titles_from_category(category_page, max_depth=1)
print(f"Trouvé {len(titles)} titres d'articles dans {CATEGORY_TITLE}.")

Découpage des articles

Une fois que nous avons les articles, nous devons les diviser en sections plus petites et semi-autonomes. C'est crucial car les modèles GPT ne peuvent traiter qu'une quantité limitée de texte à la fois. Nous écartons les sections moins pertinentes et nettoyons le texte en supprimant les balises de référence et les espaces blancs.

def clean_section(section):
    titles, text = section
    text = re.sub(r"<ref.*?</ref>", "", text)
    text = text.strip()
    return (titles, text)
 
wikipedia_sections = [clean_section(ws) for ws in wikipedia_sections]

Création des embeddings pour les sections

Avec les sections préparées, nous utilisons l'API OpenAI pour générer des embeddings pour chaque section. Cette étape transforme le texte en un format numérique qui peut être facilement recherché et comparé.

from openai import OpenAI
import pandas as pd
 
client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
 
embeddings = []
for batch_start in range(0, len(wikipedia_strings), BATCH_SIZE):
    batch_end = batch_start + BATCH_SIZE
    batch = wikipedia_strings[batch_start:batch_end]
    response = client.embeddings.create(model=EMBEDDING_MODEL, input=batch)
    embeddings.extend([e.embedding for e in response.data])
 
df = pd.DataFrame({"text": wikipedia_strings, "embedding": embeddings})

Stockage des embeddings

Enfin, nous stockons les embeddings dans un fichier CSV pour un accès et une récupération faciles. Pour des datasets plus grands, envisagez d'utiliser une base de données vectorielle pour de meilleures performances.

SAVE_PATH = "data/winter_olympics_2022.csv"
df.to_csv(SAVE_PATH, index=False)

Ce processus nous permet de préparer et de rechercher efficacement à travers de grandes collections de données textuelles, ce qui en fait un outil puissant pour des applications comme la réponse aux questions et la recherche d'informations.

Référence : Embedding Wikipedia articles for search par Ted Sanders.