Wikipedia QA Opean AI example¶

import openai
import os
import json

OPEN_AI_API_KEY = os.getenv("OPENAI_API_KEY")
openai.api_key = OPEN_AI_API_KEY
client = openai.OpenAI()

# creating a prompt
question_prompt = """
Who is the owner of twitter?
Answer: 
"""

# Use completion endpoint
completion = client.chat.completions.create(
    model="gpt-3.5-turbo",
    messages=[
        {
            "role": "system",
            "content": "You are a helpful assistant",
        },
        {
            "role": "user",
            "content": question_prompt,
        },
    ],
)

import requests
import pandas as pd
from dateutil.parser import parse, ParserError

def clean_wikipedia_data(df: pd.DataFrame) -> pd.DataFrame:
    df_cleaned = df.copy()
    df_cleaned = df = df_cleaned[
        (df_cleaned["text"].str.len() > 0) & (~df_cleaned["text"].str.startswith("=="))
    ]

    return df_cleaned


def parse_dates(df_cleaned: pd.DataFrame) -> pd.DataFrame:
    # In some cases dates are used as headings instead of being part of the
    # text sample; adjust so dated text samples start with dates
    prefix = ""
    for i, row in df_cleaned.iterrows():
        # If the row already has " - ", it already has the needed date prefix
        if " – " not in row["text"]:
            try:
                # If the row's text is a date, set it as the new prefix
                parse(row["text"])
                prefix = row["text"]
            except ParserError:
                # If the row's text isn't a date, add the prefix
                row["text"] = prefix + " – " + row["text"]
                
    df_cleaned = df_cleaned[df_cleaned["text"].str.contains(" – ")].reset_index(
        drop=True
    )
    
    return df_cleaned


df_cleaned = clean_wikipedia_data(df)
df_cleaned = parse_dates(df_cleaned)

df_cleaned.to_csv("data/wikipedia_data.csv")

EMBEDDING_MODEL_NAME = "text-embedding-ada-002"

def get_embedding(text, model="text-embedding-3-small"):
    text = text.replace("\n", " ")
    return client.embeddings.create(input=[text], model=model).data[0].embedding


# df["embeddings"] = df.text.apply(
#     lambda x: get_embedding(x, model=EMBEDDING_MODEL_NAME)
# )
# df.to_csv("data/wikipedia_embeddings.csv")

question = "Who is the owner of twitter?"
questionn_embeddings = get_embedding(question, model=EMBEDDING_MODEL_NAME)

from scipy import spatial


def distances_from_embeddings(
    query_embedding: list[float],
    embeddings: list[list[float]],
    distance_metric="cosine",
) -> list[list]:
    """Return the distances between a query embedding and a list of embeddings."""
    distance_metrics = {
        "cosine": spatial.distance.cosine,
        "L1": spatial.distance.cityblock,
        "L2": spatial.distance.euclidean,
        "Linf": spatial.distance.chebyshev,
    }
    distances = [
        distance_metrics[distance_metric](query_embedding, embedding)
        for embedding in embeddings
    ]
    return distances

df = df.assign(
    distances=distances_from_embeddings(
        questionn_embeddings, df["embeddings"], distance_metric="cosine"
    )
)

df.to_csv("data/wikipedia_embeddings_distance.csv")

import os
import pandas as pd
import openai
import numpy as np
import tiktoken

df = pd.read_csv("data/wikipedia_embeddings_distance.csv", index_col =0).sort_values("distances")

def get_number_of_tokens(text: str, tokenizer) -> pd.DataFrame:
    return len(tokenizer.encode(text))

df = df.assign(
    length_token = df['text'].apply(lambda x: get_number_of_tokens(x, tokenizer))
)

def create_prompt(question, df, tokenizer, max_token_count):
    """
    Given a question and a dataframe containing rows of text and their
    embeddings, return a text prompt to send to a Completion model
    """

    # Count the number of tokens in the prompt template and question
    prompt_template = """
    Answer the question based on the context below, and if the question
    can't be answered based on the context, say "I don't know"

    Context: 

    {}

    ---

    Question: {}
    Answer:"""

    current_token_count = len(tokenizer.encode(prompt_template)) + len(
        tokenizer.encode(question)
    )

    context = []
    for text in df.sort_values("distances")["text"].values:
        # Increase the counter based on the number of tokens in this row
        text_token_count = len(tokenizer.encode(text))
        current_token_count += text_token_count

        # Add the row of text to the list if we haven't exceeded the max
        if current_token_count <= max_token_count:
            context.append(text)
        else:
            break

    return prompt_template.format("n\n###\n\n".join(context), question)

completion = client.chat.completions.create(
    model="gpt-3.5-turbo",
    messages=[
        {
            "role": "system",
            "content": "You are a helpful assistant",
        },
        {
            "role": "user",
            "content": create_prompt(question, df, tokenizer, max_token_count),
        },
    ],
)