Code search using embeddings

Open in GithubThis notebook shows how Ada embeddings can be used to implement semantic code search.  For this demonstration, we use our own openai-python code repository.  We implement a simple version of file parsing and extracting of functions from python files, which can be embedded, indexed, and queried.Helper Functions We first setup some simple parsing functions that allow us to extract important information from our codebase.import pandas as pdfrom pathlib import PathDEF_PREFIXES = ['def ', 'async def ']NEWLINE = '\n'def get_function_name(code): ''' Extract function name from a line beginning with 'def' or 'async def'.  ''' for prefix in DEF_PREFIXES: if code.startswith(prefix): return code[len(prefix): code.index('(')]def get_until_no_space(all_lines, i): ''' Get all lines until a line outside the function definition is found.  ''' ret = [all_lines[i]] for j in range(i + 1, len(all_lines)): if len(all_lines[j]) == 0 or all_lines[j][0] in [' ', '\t', ')']: ret.append(all_lines[j]) else: break return NEWLINE.join(ret)def get_functions(filepath): ''' Get all functions in a Python file.  ''' with open(filepath, 'r') as file: all_lines = file.read().replace('\r', NEWLINE).split(NEWLINE) for i, l in enumerate(all_lines): for prefix in DEF_PREFIXES: if l.startswith(prefix): code = get_until_no_space(all_lines, i) function_name = get_function_name(code) yield { 'code': code, 'function_name': function_name, 'filepath': filepath, } breakdef extract_functions_from_repo(code_root): ''' Extract all .py functions from the repository.  ''' code_files = list(code_root.glob('**/*.py')) num_files = len(code_files) print(f'Total number of .py files: {num_files}') if num_files == 0: print('Verify openai-python repo exists and code_root is set correctly.') return None all_funcs = [ func for code_file in code_files for func in get_functions(str(code_file)) ] num_funcs = len(all_funcs) print(f'Total number of functions extracted: {num_funcs}') return all_funcsData Loading We'll first load the openai-python folder and extract the needed information using the functions we defined above.# Set user root directory to the 'openai-python' repositoryroot_dir = Path.home()# Assumes the 'openai-python' repository exists in the user's root directorycode_root = root_dir / 'openai-python'# Extract all functions from the repositoryall_funcs = extract_functions_from_repo(code_root)Now that we have our content, we can pass the data to the text-embedding-ada-002 endpoint to get back our vector embeddings.from openai.embeddings_utils import get_embeddingdf = pd.DataFrame(all_funcs)df['code_embedding'] = df['code'].apply(lambda x: get_embedding(x, engine='text-embedding-ada-002'))df['filepath'] = df['filepath'].map(lambda x: Path(x).relative_to(code_root))df.to_csv('data/code_search_openai-python.csv', index=False)df.head() .dataframe tbody tr th:only-of-type { vertical-align: middle; } .dataframe tbody tr th { vertical-align: top; } .dataframe thead th { text-align: right; } code function_name filepath code_embedding 0 def _console_log_level():\n if openai.log i...  _console_log_level openai/util.py [0.033906757831573486, -0.00418944051489234, 0...  1 def log_debug(message, **params):\n msg = l...  log_debug openai/util.py [-0.004059609025716782, 0.004895503632724285, ...  2 def log_info(message, **params):\n msg = lo...  log_info openai/util.py [0.0048639848828315735, 0.0033139237202703953,...  3 def log_warn(message, **params):\n msg = lo...  log_warn openai/util.py [0.0024026145692914724, -0.010721310041844845,...  4 def logfmt(props):\n def fmt(key, val):\n ...  logfmt openai/util.py [0.01664826273918152, 0.01730910874903202, 0.0...  Testing Let's test our endpoint with some simple queries.  If you're familiar with the openai-python repository, you'll see that we're able to easily find functions we're looking for only a simple English description.  We define a search_functions method that takes our data that contains our embeddings, a query string, and some other configuration options


Code search using embeddings

The process of searching our database works like such: We first embed our query string (code_query) with text-embedding-ada-002 .  The reasoning here is that a query string like 'a function that reverses a string' and a function like 'def reverse(string): return string[::-1]' will be very similar when embedded.  We then calculate the cosine similarity between our query string embedding and all data points in our database.  This gives a distance between each point and our query.  We finally sort all of our data points by their distance to our query string and return the number of results requested in the function parameters.  from openai.embeddings_utils import cosine_similaritydef search_functions(df, code_query, n=3, pprint=True, n_lines=7): embedding = get_embedding(code_query, engine='text-embedding-ada-002') df['similarities'] = df.code_embedding.apply(lambda x: cosine_similarity(x, embedding)) res = df.sort_values('similarities', ascending=False).head(n) if pprint: for r in res.iterrows(): print(f'{r[1].filepath}:{r[1].function_name} score={round(r[1].similarities, 3)}') print('\n'.join(r[1].code.split('\n')[:n_lines])) print('-' * 70) return resres = search_functions(df, 'fine-tuning input data validation logic', n=3)res = search_functions(df, 'find common suffix', n=2, n_lines=10)res = search_functions(df, 'Command line interface for fine-tuning', n=1, n_lines=20)


Testing

Let's test our endpoint with some simple queries. If you're familiar with the openai-python repository, you'll see that we're able to easily find functions we're looking for only a simple English description.We define a search_functions method that takes our data that contains our embeddings, a query string, and some other configuration options. The process of searching our database works like such: We first embed our query string (code_query) with text-embedding-ada-002. The reasoning here is that a query string like 'a function that reverses a string' and a function like 'def reverse(string): return string[::-1]' will be very similar when embedded. We then calculate the cosine similarity between our query string embedding and all data points in our database. This gives a distance between each point and our query. We finally sort all of our data points by their distance to our query string and return the number of results requested in the function parameters.


Helper Functions

We first setup some simple parsing functions that allow us to extract important information from our codebase.


Data Loading

We'll first load the openai-python folder and extract the needed information using the functions we defined above.


Customizing embeddings

Open in GithubThis notebook demonstrates one way to customize OpenAI embeddings to a particular task. The input is training data in the form of [text_1, text_2, label] where label is +1 if the pairs are similar and -1 if the pairs are dissimilar. The output is a matrix that you can use to multiply your embeddings. The product of this multiplication is a 'custom embedding' that will better emphasize aspects of the text relevant to your use case. In binary classification use cases, we've seen error rates drop by as much as 50%. In the following example, I use 1,000 sentence pairs picked from the SNLI corpus. Each pair of sentences are logically entailed (i.e., one implies the other). These pairs are our positives (label = 1). We generate synthetic negatives by combining sentences from different pairs, which are presumed to not be logically entailed (label = -1). For a clustering use case, you can generate positives by creating pairs from texts in the same clusters and generate negatives by creating pairs from sentences in different clusters. With other data sets, we have seen decent improvement with as little as ~100 training examples. Of course, performance will be better with more examples.


0. Imports

# importsfrom typing import List, Tuple # for type hintsimport numpy as np # for manipulating arraysimport pandas as pd # for manipulating data in dataframesimport pickle # for saving the embeddings cacheimport plotly.express as px # for plotsimport random # for generating run IDsfrom sklearn.model_selection import train_test_split # for splitting train & test dataimport torch # for matrix optimizationfrom openai.embeddings_utils import get_embedding, cosine_similarity # for embeddings1. Inputs Most inputs are here. The key things to change are where to load your datset from, where to save a cache of embeddings to, and which embedding engine you want to use. Depending on how your data is formatted, you'll want to rewrite the process_input_data function.# input parametersembedding_cache_path = 'data/snli_embedding_cache.pkl' # embeddings will be saved/loaded heredefault_embedding_engine = 'babbage-similarity' # text-embedding-ada-002 is recommendednum_pairs_to_embed = 1000 # 1000 is arbitrarylocal_dataset_path = 'data/snli_1.0_train_2k.csv' # download from: https://nlp.stanford.edu/projects/snli/def process_input_data(df: pd.DataFrame) -> pd.DataFrame: # you can customize this to preprocess your own dataset # output should be a dataframe with 3 columns: text_1, text_2, label (1 for similar, -1 for dissimilar) df['label'] = df['gold_label'] df = df[df['label'].isin(['entailment'])] df['label'] = df['label'].apply(lambda x: {'entailment': 1, 'contradiction': -1}[x]) df = df.rename(columns={'sentence1': 'text_1', 'sentence2': 'text_2'}) df = df[['text_1', 'text_2', 'label']] df = df.head(num_pairs_to_embed) return df


2. Load and process input data

# load datadf = pd.read_csv(local_dataset_path)# process input datadf = process_input_data(df) # this demonstrates training data containing only positives# view datadf.head() .dataframe tbody tr th:only-of-type { vertical-align: middle; } .dataframe tbody tr th { vertical-align: top; } .dataframe thead th { text-align: right; } text_1 text_2 label 2 A person on a horse jumps over a broken down a...  A person is outdoors, on a horse.  1 4 Children smiling and waving at camera There are children present 1 7 A boy is jumping on skateboard in the middle o...  The boy does a skateboarding trick.  1 14 Two blond women are hugging one another.  There are women showing affection.  1 17 A few people in a restaurant setting, one of t...  The diners are at a restaurant.  1 3.  Split data into training test sets Note that it's important to split data into training and test sets before generating synethetic negatives or positives.  You don't want any text strings in the training data to show up in the test data.  If there's contamination, the test metrics will look better than they'll actually be in production.# split data into train and test setstest_fraction = 0.5 # 0.5 is fairly arbitraryrandom_seed = 123 # random seed is arbitrary, but is helpful in reproducibilitytrain_df, test_df = train_test_split( df, test_size=test_fraction, stratify=df['label'], random_state=random_seed)train_df.loc[:, 'dataset'] = 'train'test_df.loc[:, 'dataset'] = 'test'4.  Generate synthetic negatives This is another piece of the code that you will need to modify to match your use case.  If you have data with positives and negatives, you can skip this section.  If you have data with only positives, you can mostly keep it as is, where it generates negatives only.  If you have multiclass data, you will want to generate both positives and negatives.  The positives can be pairs of text that share labels, and the negatives can be pairs of text that do not share labels.  The final output should be a dataframe with text pairs, where each pair is labeled -1 or 1.# generate negativesdef dataframe_of_negatives(dataframe_of_positives: pd.DataFrame) -> pd.DataFrame: '''Return dataframe of negative pairs made by combining elements of positive pairs.''' texts = set(dataframe_of_positives['text_1'].values) | set( dataframe_of_positives['text_2'].values ) all_pairs = {(t1, t2) for t1 in texts for t2 in texts if t1 < t2} positive_pairs = set( tuple(text_pair) for text_pair in dataframe_of_positives[['text_1', 'text_2']].values ) negative_pairs = all_pairs - positive_pairs df_of_negatives = pd.DataFrame(list(negative_pairs), columns=['text_1', 'text_2']) df_of_negatives['label'] = -1 return df_of_negativesnegatives_per_positive = ( 1 # it will work at higher values too, but more data will be slower)# generate negatives for training datasettrain_df_negatives = dataframe_of_negatives(train_df)train_df_negatives['dataset'] = 'train'# generate negatives for test datasettest_df_negatives = dataframe_of_negatives(test_df)test_df_negatives['dataset'] = 'test'# sample negatives and combine with positivestrain_df = pd.concat( [ train_df, train_df_negatives.sample( n=len(train_df) * negatives_per_positive, random_state=random_seed ), ])test_df = pd.concat( [ test_df, test_df_negatives.sample( n=len(test_df) * negatives_per_positive, random_state=random_seed ), ])df = pd.concat([train_df, test_df])5.  Calculate embeddings and cosine similarities Here, I create a cache to save the embeddings


1. Inputs

Most inputs are here. The key things to change are where to load your datset from, where to save a cache of embeddings to, and which embedding engine you want to use.Depending on how your data is formatted, you'll want to rewrite the process_input_data function.


3. Split data into training test sets

Note that it's important to split data into training and test sets before generating synethetic negatives or positives. You don't want any text strings in the training data to show up in the test data. If there's contamination, the test metrics will look better than they'll actually be in production.


2. Load and process input data

This is handy so that you don't have to pay again if you want to run the code again.# establish a cache of embeddings to avoid recomputing# cache is a dict of tuples (text, engine) -> embeddingtry: with open(embedding_cache_path, 'rb') as f: embedding_cache = pickle.load(f)except FileNotFoundError: precomputed_embedding_cache_path = 'https://cdn.openai.com/API/examples/data/snli_embedding_cache.pkl' embedding_cache = pd.read_pickle(precomputed_embedding_cache_path)# this function will get embeddings from the cache and save them there afterwarddef get_embedding_with_cache( text: str, engine: str = default_embedding_engine, embedding_cache: dict = embedding_cache, embedding_cache_path: str = embedding_cache_path,) -> list: if (text, engine) not in embedding_cache.keys(): # if not in cache, call API to get embedding embedding_cache[(text, engine)] = get_embedding(text, engine) # save embeddings cache to disk after each update with open(embedding_cache_path, 'wb') as embedding_cache_file: pickle.dump(embedding_cache, embedding_cache_file) return embedding_cache[(text, engine)]# create column of embeddingsfor column in ['text_1', 'text_2']: df[f'{column}_embedding'] = df[column].apply(get_embedding_with_cache)# create column of cosine similarity between embeddingsdf['cosine_similarity'] = df.apply( lambda row: cosine_similarity(row['text_1_embedding'], row['text_2_embedding']), axis=1,)6 .  Plot distribution of cosine similarity Here we measure similarity of text using cosine similarity.  In our experience, most distance functions (L1, L2, cosine similarity) all work about the same.  Note that our embeddings are already normalized to length 1, so cosine similarity is equivalent to dot product.  The graphs show how much the overlap there is between the distribution of cosine similarities for similar and dissimilar pairs.  If there is a high amount of overlap, that means there are some dissimilar pairs with greater cosine similarity than some similar pairs.  The accuracy I compute is the accuracy of a simple rule that predicts 'similar (1)' if the cosine similarity is above some threshold X and otherwise predicts 'dissimilar (0)'.# calculate accuracy (and its standard error) of predicting label=1 if similarity>x# x is optimized by sweeping from -1 to 1 in steps of 0.01def accuracy_and_se(cosine_similarity: float, labeled_similarity: int) -> Tuple[float]: accuracies = [] for threshold_thousandths in range(-1000, 1000, 1): threshold = threshold_thousandths / 1000 total = 0 correct = 0 for cs, ls in zip(cosine_similarity, labeled_similarity): total += 1 if cs > threshold: prediction = 1 else: prediction = -1 if prediction == ls: correct += 1 accuracy = correct / total accuracies.append(accuracy) a = max(accuracies) n = len(cosine_similarity) standard_error = (a * (1 - a) / n) ** 0.5 # standard error of binomial return a, standard_error# check that training and test sets are balancedpx.histogram( df, x='cosine_similarity', color='label', barmode='overlay', width=500, facet_row='dataset',).show()for dataset in ['train', 'test']: data = df[df['dataset'] == dataset] a, se = accuracy_and_se(data['cosine_similarity'], data['label']) print(f'{dataset} accuracy: {a:0.1%} ± {1.96 * se:0.1%}')


4. Generate synthetic negatives

This is another piece of the code that you will need to modify to match your use case.If you have data with positives and negatives, you can skip this section.If you have data with only positives, you can mostly keep it as is, where it generates negatives only.If you have multiclass data, you will want to generate both positives and negatives. The positives can be pairs of text that share labels, and the negatives can be pairs of text that do not share labels.The final output should be a dataframe with text pairs, where each pair is labeled -1 or 1.


6. Plot distribution of cosine similarity

Here we measure similarity of text using cosine similarity. In our experience, most distance functions (L1, L2, cosine similarity) all work about the same. Note that our embeddings are already normalized to length 1, so cosine similarity is equivalent to dot product.The graphs show how much the overlap there is between the distribution of cosine similarities for similar and dissimilar pairs. If there is a high amount of overlap, that means there are some dissimilar pairs with greater cosine similarity than some similar pairs.The accuracy I compute is the accuracy of a simple rule that predicts 'similar (1)' if the cosine similarity is above some threshold X and otherwise predicts 'dissimilar (0)'.


5. Calculate embeddings and cosine similarities

Here, I create a cache to save the embeddings. This is handy so that you don't have to pay again if you want to run the code again.


7. Optimize the matrix using the training data provided

def embedding_multiplied_by_matrix( embedding: List[float], matrix: torch.tensor) -> np.array: embedding_tensor = torch.tensor(embedding).float() modified_embedding = embedding_tensor @ matrix modified_embedding = modified_embedding.detach().numpy() return modified_embedding# compute custom embeddings and new cosine similaritiesdef apply_matrix_to_embeddings_dataframe(matrix: torch.tensor, df: pd.DataFrame): for column in ['text_1_embedding', 'text_2_embedding']: df[f'{column}_custom'] = df[column].apply( lambda x: embedding_multiplied_by_matrix(x, matrix) ) df['cosine_similarity_custom'] = df.apply( lambda row: cosine_similarity( row['text_1_embedding_custom'], row['text_2_embedding_custom'] ), axis=1, )def optimize_matrix( modified_embedding_length: int = 2048, # in my brief experimentation, bigger was better (2048 is length of babbage encoding) batch_size: int = 100, max_epochs: int = 100, learning_rate: float = 100.0, # seemed to work best when similar to batch size - feel free to try a range of values dropout_fraction: float = 0.0, # in my testing, dropout helped by a couple percentage points (definitely not necessary) df: pd.DataFrame = df, print_progress: bool = True, save_results: bool = True,) -> torch.tensor: '''Return matrix optimized to minimize loss on training data.''' run_id = random.randint(0, 2 ** 31 - 1) # (range is arbitrary) # convert from dataframe to torch tensors # e is for embedding, s for similarity label def tensors_from_dataframe( df: pd.DataFrame, embedding_column_1: str, embedding_column_2: str, similarity_label_column: str, ) -> Tuple[torch.tensor]: e1 = np.stack(np.array(df[embedding_column_1].values)) e2 = np.stack(np.array(df[embedding_column_2].values)) s = np.stack(np.array(df[similarity_label_column].astype('float').values)) e1 = torch.from_numpy(e1).float() e2 = torch.from_numpy(e2).float() s = torch.from_numpy(s).float() return e1, e2, s e1_train, e2_train, s_train = tensors_from_dataframe( df[df['dataset'] == 'train'], 'text_1_embedding', 'text_2_embedding', 'label' ) e1_test, e2_test, s_test = tensors_from_dataframe( df[df['dataset'] == 'test'], 'text_1_embedding', 'text_2_embedding', 'label' ) # create dataset and loader dataset = torch.utils.data.TensorDataset(e1_train, e2_train, s_train) train_loader = torch.utils.data.DataLoader( dataset, batch_size=batch_size, shuffle=True ) # define model (similarity of projected embeddings) def model(embedding_1, embedding_2, matrix, dropout_fraction=dropout_fraction): e1 = torch.nn.functional.dropout(embedding_1, p=dropout_fraction) e2 = torch.nn.functional.dropout(embedding_2, p=dropout_fraction) modified_embedding_1 = e1 @ matrix # @ is matrix multiplication modified_embedding_2 = e2 @ matrix similarity = torch.nn.functional.cosine_similarity( modified_embedding_1, modified_embedding_2 ) return similarity # define loss function to minimize def mse_loss(predictions, targets): difference = predictions - targets return torch.sum(difference * difference) / difference.numel() # initialize projection matrix embedding_length = len(df['text_1_embedding'].values[0]) matrix = torch.randn( embedding_length, modified_embedding_length, requires_grad=True ) epochs, types, losses, accuracies, matrices = [], [], [], [], [] for epoch in range(1, 1 + max_epochs): # iterate through training dataloader for a, b, actual_similarity in train_loader: # generate prediction predicted_similarity = model(a, b, matrix) # get loss and perform backpropagation loss = mse_loss(predicted_similarity, actual_similarity) loss.backward() # update the weights with torch.no_grad(): matrix -= matrix.grad * learning_rate # set gradients to zero matrix.grad.zero_() # calculate test loss test_predictions = model(e1_test, e2_test, matrix) test_loss = mse_loss(test_predictions, s_test) # compute custom embeddings and new cosine similarities apply_matrix_to_embeddings_dataframe(matrix, df) # calculate test accuracy for dataset in ['train', 'test']: data = df[df['dataset'] == dataset] a, se = accuracy_and_se(data['cosine_similarity_custom'], data['label']) # record results of each epoch epochs.append(epoch) types.append(dataset) losses.append(loss.item() if dataset == 'train' else test_loss.item()) accuracies.append(a) matrices.append(matrix.detach().numpy()) # optionally print accuracies if print_progress is True: print( f'Epoch {epoch}/{max_epochs}: {dataset} accuracy: {a:0.1%} ± {1.96 * se:0.1%}' ) data = pd.DataFrame( {'epoch': epochs, 'type': types, 'loss': losses, 'accuracy': accuracies} ) data['run_id'] = run_id data['modified_embedding_length'] = modified_embedding_length data['batch_size'] = batch_size data['max_epochs'] = max_epochs data['learning_rate'] = learning_rate data['dropout_fraction'] = dropout_fraction data[ 'matrix' ] = matrices # saving every single matrix can get big; feel free to delete/change if save_results is True: data.to_csv(f'{run_id}_optimization_results.csv', index=False) return data# example hyperparameter search# I recommend starting with max_epochs=10 while initially exploringresults = []max_epochs = 30dropout_fraction = 0.2for batch_size, learning_rate in [(10, 10), (100, 100), (1000, 1000)]: result = optimize_matrix( batch_size=batch_size, learning_rate=learning_rate, max_epochs=max_epochs, dropout_fraction=dropout_fraction, save_results=False, ) results.append(result)runs_df = pd.concat(results)# plot training loss and test loss over timepx.line( runs_df, line_group='run_id', x='epoch', y='loss', color='type', hover_data=['batch_size', 'learning_rate', 'dropout_fraction'], facet_row='learning_rate', facet_col='batch_size', width=500,).show()# plot accuracy over timepx.line( runs_df, line_group='run_id', x='epoch', y='accuracy', color='type', hover_data=['batch_size', 'learning_rate', 'dropout_fraction'], facet_row='learning_rate', facet_col='batch_size', width=500,).show()8


8. Plot the before & after, showing the results of the best matrix found during training

The better the matrix is, the more cleanly it will separate the similar and dissimilar pairs.


Plot the before & after, showing the results of the best matrix found during training The better the matrix is, the more cleanly it will separate the similar and dissimilar pairs.# apply result of best run to original databest_run = runs_df.sort_values(by='accuracy', ascending=False).iloc[0]best_matrix = best_run['matrix']apply_matrix_to_embeddings_dataframe(best_matrix, df)# plot similarity distribution BEFORE customizationpx.histogram( df, x='cosine_similarity', color='label', barmode='overlay', width=500, facet_row='dataset',).show()test_df = df[df['dataset'] == 'test']a, se = accuracy_and_se(test_df['cosine_similarity'], test_df['label'])print(f'Test accuracy: {a:0.1%} ± {1.96 * se:0.1%}')# plot similarity distribution AFTER customizationpx.histogram( df, x='cosine_similarity_custom', color='label', barmode='overlay', width=500, facet_row='dataset',).show()a, se = accuracy_and_se(test_df['cosine_similarity_custom'], test_df['label'])print(f'Test accuracy after customization: {a:0.1%} ± {1.96 * se:0.1%}')best_matrix # this is what you can multiply your embeddings by


Fine tuning classification example

Open in GithubWe will fine-tune an ada classifier to distinguish between the two sports: Baseball and Hockey.from sklearn.datasets import fetch_20newsgroupsimport pandas as pdimport openaicategories = ['rec.sport.baseball', 'rec.sport.hockey']sports_dataset = fetch_20newsgroups(subset='train', shuffle=True, random_state=42, categories=categories)Data exploration The newsgroup dataset can be loaded using sklearn.  First we will look at the data itself:print(sports_dataset['data'][0])sports_dataset.target_names[sports_dataset['target'][0]]len_all, len_baseball, len_hockey = len(sports_dataset.data), len([e for e in sports_dataset.target if e == 0]), len([e for e in sports_dataset.target if e == 1])print(f'Total examples: {len_all}, Baseball examples: {len_baseball}, Hockey examples: {len_hockey}')One sample from the baseball category can be seen above.  It is an email to a mailing list.  We can observe that we have 1197 examples in total, which are evenly split between the two sports.Data Preparation We transform the dataset into a pandas dataframe, with a column for prompt and completion.  The prompt contains the email from the mailing list, and the completion is a name of the sport, either hockey or baseball.  For demonstration purposes only and speed of fine-tuning we take only 300 examples.  In a real use case the more examples the better the performance.import pandas as pdlabels = [sports_dataset.target_names[x].split('.')[-1] for x in sports_dataset['target']]texts = [text.strip() for text in sports_dataset['data']]df = pd.DataFrame(zip(texts, labels), columns = ['prompt','completion']) #[:300]df.head() .dataframe tbody tr th:only-of-type { vertical-align: middle; } .dataframe tbody tr th { vertical-align: top; } .dataframe thead th { text-align: right; } prompt completion 0 From: dougb@comm.mot.com (Doug Bank)\nSubject:...  baseball 1 From: gld@cunixb.cc.columbia.edu (Gary L Dare)...  hockey 2 From: rudy@netcom.com (Rudy Wade)\nSubject: Re...  baseball 3 From: monack@helium.gas.uug.arizona.edu (david...  hockey 4 Subject: Let it be Known\nFrom: <ISSBTL@BYUVM....  baseball Both baseball and hockey are single tokens.  We save the dataset as a jsonl file.df.to_json('sport2.jsonl', orient='records', lines=True)Data Preparation tool We can now use a data preparation tool which will suggest a few improvements to our dataset before fine-tuning.  Before launching the tool we update the openai library to ensure we're using the latest data preparation tool.  We additionally specify -q which auto-accepts all suggestions.!pip install --upgrade openai!openai tools fine_tunes.prepare_data -f sport2.jsonl -qThe tool helpfully suggests a few improvements to the dataset and splits the dataset into training and validation set.  A suffix between a prompt and a completion is necessary to tell the model that the input text has stopped, and that it now needs to predict the class.  Since we use the same separator in each example, the model is able to learn that it is meant to predict either baseball or hockey following the separator.  A whitespace prefix in completions is useful, as most word tokens are tokenized with a space prefix.  The tool also recognized that this is likely a classification task, so it suggested to split the dataset into training and validation datasets.  This will allow us to easily measure expected performance on new data.Fine-tuning The tool suggests we run the following command to train the dataset


Fine tuning classification example

Since this is a classification task, we would like to know what the generalization performance on the provided validation set is for our classification use case .  The tool suggests to add --compute_classification_metrics --classification_positive_class ' baseball' in order to compute the classification metrics.  We can simply copy the suggested command from the CLI tool.  We specifically add -m ada to fine-tune a cheaper and faster ada model, which is usually comperable in performance to slower and more expensive models on classification use cases.!openai api fine_tunes.create -t 'sport2_prepared_train.jsonl' -v 'sport2_prepared_valid.jsonl' --compute_classification_metrics --classification_positive_class ' baseball' -m adaThe model is successfully trained in about ten minutes.  We can see the model name is ada:ft-openai-2021-07-30-12-26-20, which we can use for doing inference.[Advanced] Results and expected model performance We can now download the results file to observe the expected performance on a held out validation set.!openai api fine_tunes.results -i ft-2zaA7qi0rxJduWQpdvOvmGn3 > result.csvresults = pd.read_csv('result.csv')results[results['classification/accuracy'].notnull()].tail(1) .dataframe tbody tr th:only-of-type { vertical-align: middle; } .dataframe tbody tr th { vertical-align: top; } .dataframe thead th { text-align: right; } step elapsed_tokens elapsed_examples training_loss training_sequence_accuracy training_token_accuracy classification/accuracy classification/precision classification/recall classification/auroc classification/auprc classification/f1.0 validation_loss validation_sequence_accuracy validation_token_accuracy 929 930 3027688 3720 0.044408 1.0 1.0 0.991597 0.983471 1.0 1.0 1.0 0.991667 NaN NaN NaN The accuracy reaches 99.6%.  On the plot below we can see how accuracy on the validation set increases during the training run.results[results['classification/accuracy'].notnull()]['classification/accuracy'].plot()Using the model We can now call the model to get the predictions.test = pd.read_json('sport2_prepared_valid.jsonl', lines=True)test.head() .dataframe tbody tr th:only-of-type { vertical-align: middle; } .dataframe tbody tr th { vertical-align: top; } .dataframe thead th { text-align: right; } prompt completion 0 From: gld@cunixb.cc.columbia.edu (Gary L Dare)...  hockey 1 From: smorris@venus.lerc.nasa.gov (Ron Morris ...  hockey 2 From: golchowy@alchemy.chem.utoronto.ca (Geral...  hockey 3 From: krattige@hpcc01.corp.hp.com (Kim Krattig...  baseball 4 From: warped@cs.montana.edu (Doug Dolven)\nSub...  baseball We need to use the same separator following the prompt which we used during fine-tuning.  In this case it is \n\n###\n\n.  Since we're concerned with classification, we want the temperature to be as low as possible, and we only require one token completion to determine the prediction of the model.ft_model = 'ada:ft-openai-2021-07-30-12-26-20'res = openai.Completion.create(model=ft_model, prompt=test['prompt'][0] + '\n\n###\n\n', max_tokens=1, temperature=0)res['choices'][0]['text']To get the log probabilities, we can specify logprobs parameter on the completion requestres = openai.Completion.create(model=ft_model, prompt=test['prompt'][0] + '\n\n###\n\n', max_tokens=1, temperature=0, logprobs=2)res['choices'][0]['logprobs']['top_logprobs'][0]We can see that the model predicts hockey as a lot more likely than baseball, which is the correct prediction.  By requesting log_probs, we can see the prediction (log) probability for each class.Generalization Interestingly, our fine-tuned classifier is quite versatile


Fine tuning classification example

Despite being trained on emails to different mailing lists, it also successfully predicts tweets.sample_hockey_tweet = '''Thank you to the @Canes and all you amazing Caniacs that have been so supportive! You guys are some of the best fans in the NHL without a doubt! Really excited to start this new chapter in my career with the @DetroitRedWings !!'''res = openai.Completion.create(model=ft_model, prompt=sample_hockey_tweet + '\n\n###\n\n', max_tokens=1, temperature=0, logprobs=2)res['choices'][0]['text']sample_baseball_tweet='''BREAKING: The Tampa Bay Rays are finalizing a deal to acquire slugger Nelson Cruz from the Minnesota Twins, sources tell ESPN.'''res = openai.Completion.create(model=ft_model, prompt=sample_baseball_tweet + '\n\n###\n\n', max_tokens=1, temperature=0, logprobs=2)res['choices'][0]['text']


Data Preparation

We transform the dataset into a pandas dataframe, with a column for prompt and completion. The prompt contains the email from the mailing list, and the completion is a name of the sport, either hockey or baseball. For demonstration purposes only and speed of fine-tuning we take only 300 examples. In a real use case the more examples the better the performance.


Data exploration

The newsgroup dataset can be loaded using sklearn. First we will look at the data itself:


Fine-tuning

The tool suggests we run the following command to train the dataset. Since this is a classification task, we would like to know what the generalization performance on the provided validation set is for our classification use case. The tool suggests to add --compute_classification_metrics --classification_positive_class ' baseball' in order to compute the classification metrics.We can simply copy the suggested command from the CLI tool. We specifically add -m ada to fine-tune a cheaper and faster ada model, which is usually comperable in performance to slower and more expensive models on classification use cases.


Using the model

We can now call the model to get the predictions.


Generalization

Interestingly, our fine-tuned classifier is quite versatile. Despite being trained on emails to different mailing lists, it also successfully predicts tweets.


[Advanced] Results and expected model performance

We can now download the results file to observe the expected performance on a held out validation set.


Data Preparation tool

We can now use a data preparation tool which will suggest a few improvements to our dataset before fine-tuning. Before launching the tool we update the openai library to ensure we're using the latest data preparation tool. We additionally specify -q which auto-accepts all suggestions.


Using embeddings

Open in GithubThis notebook contains some helpful snippets you can use to embed text with the 'text-embedding-ada-002' model via the OpenAI API.import openaiembedding = openai.Embedding.create( input='Your text goes here', model='text-embedding-ada-002')['data'][0]['embedding']len(embedding)It's recommended to use the 'tenacity' package or another exponential backoff implementation to better manage API rate limits, as hitting the API too much too fast can trigger rate limits. Using the following function ensures you get your embeddings as fast as possible.# Negative example (slow and rate-limited)import openainum_embeddings = 10000 # Some large numberfor i in range(num_embeddings): embedding = openai.Embedding.create( input='Your text goes here', model='text-embedding-ada-002' )['data'][0]['embedding'] print(len(embedding))# Best practiceimport openaifrom tenacity import retry, wait_random_exponential, stop_after_attempt# Retry up to 6 times with exponential backoff, starting at 1 second and maxing out at 20 seconds delay@retry(wait=wait_random_exponential(min=1, max=20), stop=stop_after_attempt(6))def get_embedding(text: str, model='text-embedding-ada-002') -> list[float]: return openai.Embedding.create(input=[text], model=model)['data'][0]['embedding']embedding = get_embedding('Your text goes here', model='text-embedding-ada-002')print(len(embedding))


Get embeddings from dataset

Open in GithubThis notebook gives an example on how to get embeddings from a large dataset. 1. Load the dataset The dataset used in this example is fine-food reviews from Amazon. The dataset contains a total of 568,454 food reviews Amazon users left up to October 2012. We will use a subset of this dataset, consisting of 1,000 most recent reviews for illustration purposes. The reviews are in English and tend to be positive or negative. Each review has a ProductId, UserId, Score, review title (Summary) and review body (Text). We will combine the review summary and review text into a single combined text. The model will encode this combined text and it will output a single vector embedding.To run this notebook, you will need to install: pandas, openai, transformers, plotly, matplotlib, scikit-learn, torch (transformer dep), torchvision, and scipy.# importsimport pandas as pdimport tiktokenfrom openai.embeddings_utils import get_embedding# embedding model parametersembedding_model = 'text-embedding-ada-002'embedding_encoding = 'cl100k_base' # this the encoding for text-embedding-ada-002max_tokens = 8000 # the maximum for text-embedding-ada-002 is 8191# load & inspect datasetinput_datapath = 'data/fine_food_reviews_1k.csv' # to save space, we provide a pre-filtered datasetdf = pd.read_csv(input_datapath, index_col=0)df = df[['Time', 'ProductId', 'UserId', 'Score', 'Summary', 'Text']]df = df.dropna()df['combined'] = ( 'Title: ' + df.Summary.str.strip() + '; Content: ' + df.Text.str.strip())df.head(2) .dataframe tbody tr th:only-of-type { vertical-align: middle; } .dataframe tbody tr th { vertical-align: top; } .dataframe thead th { text-align: right; } Time ProductId UserId Score Summary Text combined 0 1351123200 B003XPF9BO A3R7JR3FMEBXQB 5 where does one start...and stop... with a tre... Wanted to save some to bring to my Chicago fam... Title: where does one start...and stop... wit... 1 1351123200 B003JK537S A3JBPC3WFUT5ZP 1 Arrived in pieces Not pleased at all. When I opened the box, mos... Title: Arrived in pieces; Content: Not pleased... # subsample to 1k most recent reviews and remove samples that are too longtop_n = 1000df = df.sort_values('Time').tail(top_n * 2) # first cut to first 2k entries, assuming less than half will be filtered outdf.drop('Time', axis=1, inplace=True)encoding = tiktoken.get_encoding(embedding_encoding)# omit reviews that are too long to embeddf['n_tokens'] = df.combined.apply(lambda x: len(encoding.encode(x)))df = df[df.n_tokens <= max_tokens].tail(top_n)len(df)


1. Load the dataset

The dataset used in this example is fine-food reviews from Amazon. The dataset contains a total of 568,454 food reviews Amazon users left up to October 2012. We will use a subset of this dataset, consisting of 1,000 most recent reviews for illustration purposes. The reviews are in English and tend to be positive or negative. Each review has a ProductId, UserId, Score, review title (Summary) and review body (Text).We will combine the review summary and review text into a single combined text. The model will encode this combined text and it will output a single vector embedding.


2. Get embeddings and save them for future reuse

# Ensure you have your API key set in your environment per the README: https://github.com/openai/openai-python#usage# This may take a few minutesdf['embedding'] = df.combined.apply(lambda x: get_embedding(x, engine=embedding_model))df.to_csv('data/fine_food_reviews_with_embeddings_1k.csv')


Recommendation using embeddings and nearest neighbor search

The cache is saved as a Python pickle file.# establish a cache of embeddings to avoid recomputing# cache is a dict of tuples (text, model) -> embedding, saved as a pickle file# set path to embedding cacheembedding_cache_path = 'data/recommendations_embeddings_cache.pkl'# load the cache if it exists, and save a copy to disktry: embedding_cache = pd.read_pickle(embedding_cache_path)except FileNotFoundError: embedding_cache = {}with open(embedding_cache_path, 'wb') as embedding_cache_file: pickle.dump(embedding_cache, embedding_cache_file)# define a function to retrieve embeddings from the cache if present, and otherwise request via the APIdef embedding_from_string( string: str, model: str = EMBEDDING_MODEL, embedding_cache=embedding_cache) -> list: '''Return embedding of given string, using a cache to avoid recomputing.''' if (string, model) not in embedding_cache.keys(): embedding_cache[(string, model)] = get_embedding(string, model) with open(embedding_cache_path, 'wb') as embedding_cache_file: pickle.dump(embedding_cache, embedding_cache_file) return embedding_cache[(string, model)]Let's check that it works by getting an embedding.# as an example, take the first description from the datasetexample_string = df['description'].values[0]print(f'\nExample string: {example_string}')# print the first 10 dimensions of the embeddingexample_embedding = embedding_from_string(example_string)print(f'\nExample embedding: {example_embedding[:10]}...')4 .  Recommend similar articles based on embeddings To find similar articles, let's follow a three-step plan: Get the similarity embeddings of all the article descriptions Calculate the distance between a source title and all other articles Print out the other articles closest to the source title def print_recommendations_from_strings( strings: list[str], index_of_source_string: int, k_nearest_neighbors: int = 1, model=EMBEDDING_MODEL,) -> list[int]: '''Print out the k nearest neighbors of a given string.''' # get embeddings for all strings embeddings = [embedding_from_string(string, model=model) for string in strings] # get the embedding of the source string query_embedding = embeddings[index_of_source_string] # get distances between the source embedding and other embeddings (function from embeddings_utils.py) distances = distances_from_embeddings(query_embedding, embeddings, distance_metric='cosine') # get indices of nearest neighbors (function from embeddings_utils.py) indices_of_nearest_neighbors = indices_of_nearest_neighbors_from_distances(distances) # print out source string query_string = strings[index_of_source_string] print(f'Source string: {query_string}') # print out its k nearest neighbors k_counter = 0 for i in indices_of_nearest_neighbors: # skip any strings that are identical matches to the starting string if query_string == strings[i]: continue # stop after printing out k articles if k_counter >= k_nearest_neighbors: break k_counter += 1 # print out the similar strings and their distances print( f''' --- Recommendation #{k_counter} (nearest neighbor {k_counter} of {k_nearest_neighbors}) --- String: {strings[i]} Distance: {distances[i]:0.3f}''' ) return indices_of_nearest_neighbors5


Open in GithubRecommendations are widespread across the web.  'Bought that item? Try these similar items.' 'Enjoy that book? Try these similar titles.' 'Not the help page you were looking for? Try these similar pages.' This notebook demonstrates how to use embeddings to find similar items to recommend.  In particular, we use AG's corpus of news articles as our dataset.  Our model will answer the question: given an article, what other articles are most similar to it?1.  Imports First, let's import the packages and functions we'll need for later.  If you don't have these, you'll need to install them.  You can install them via your terminal by running pip install {package_name}, e.g.  pip install pandas.# importsimport pandas as pdimport picklefrom openai.embeddings_utils import ( get_embedding, distances_from_embeddings, tsne_components_from_embeddings, chart_from_components, indices_of_nearest_neighbors_from_distances,)# constantsEMBEDDING_MODEL = 'text-embedding-ada-002'2.  Load data Next, let's load the AG news data and see what it looks like.# load data (full dataset available at http://groups.di.unipi.it/~gulli/AG_corpus_of_news_articles.html)dataset_path = 'data/AG_news_samples.csv'df = pd.read_csv(dataset_path)# print dataframen_examples = 5df.head(n_examples) .dataframe tbody tr th:only-of-type { vertical-align: middle; } .dataframe tbody tr th { vertical-align: top; } .dataframe thead th { text-align: right; } title description label_int label 0 World Briefings BRITAIN: BLAIR WARNS OF CLIMATE THREAT Prime M...  1 World 1 Nvidia Puts a Firewall on a Motherboard (PC Wo...  PC World - Upcoming chip set will include buil...  4 Sci/Tech 2 Olympic joy in Greek, Chinese press Newspapers in Greece reflect a mixture of exhi...  2 Sports 3 U2 Can iPod with Pictures SAN JOSE, Calif.  -- Apple Computer (Quote, Cha...  4 Sci/Tech 4 The Dream Factory Any product, any shape, any size -- manufactur...  4 Sci/Tech Let's take a look at those same examples, but not truncated by ellipses.# print the title, description, and label of each examplefor idx, row in df.head(n_examples).iterrows(): print('') print(f'Title: {row['title']}') print(f'Description: {row['description']}') print(f'Label: {row['label']}')3.  Build cache to save embeddings Before getting embeddings for these articles, let's set up a cache to save the embeddings we generate.  In general, it's a good idea to save your embeddings so you can re-use them later.  If you don't save them, you'll pay again each time you compute them again.  The cache is a dictionary that maps tuples of (text, model) to an embedding, which is a list of floats


1. Imports

First, let's import the packages and functions we'll need for later. If you don't have these, you'll need to install them. You can install them via your terminal by running pip install {package_name}, e.g. pip install pandas.


Example recommendations Let's look for articles similar to first one, which was about Tony Blair.article_descriptions = df['description'].tolist()tony_blair_articles = print_recommendations_from_strings( strings=article_descriptions, # let's base similarity off of the article description index_of_source_string=0, # let's look at articles similar to the first one about Tony Blair k_nearest_neighbors=5, # let's look at the 5 most similar articles)Pretty good! 4 of the 5 recommendations explicitly mention Tony Blair and the fifth is an article from London about climate change, topics that might be often associated with Tony Blair.Let's see how our recommender does on the second example article about NVIDIA's new chipset with more security.chipset_security_articles = print_recommendations_from_strings( strings=article_descriptions, # let's base similarity off of the article description index_of_source_string=1, # let's look at articles similar to the second one about a more secure chipset k_nearest_neighbors=5, # let's look at the 5 most similar articles)From the printed distances, you can see that the #1 recommendation is much closer than all the others (0.11 vs 0.14+) .  And the #1 recommendation looks very similar to the starting article - it's another article from PC World about increasing computer security.  Pretty good!Appendix: Using embeddings in more sophisticated recommenders A more sophisticated way to build a recommender system is to train a machine learning model that takes in tens or hundreds of signals, such as item popularity or user click data.  Even in this system, embeddings can be a very useful signal into the recommender, especially for items that are being 'cold started' with no user data yet (e.g., a brand new product added to the catalog without any clicks yet).


2. Load data

Next, let's load the AG news data and see what it looks like.


3. Build cache to save embeddings

Before getting embeddings for these articles, let's set up a cache to save the embeddings we generate. In general, it's a good idea to save your embeddings so you can re-use them later. If you don't save them, you'll pay again each time you compute them again.The cache is a dictionary that maps tuples of (text, model) to an embedding, which is a list of floats. The cache is saved as a Python pickle file.


4. Recommend similar articles based on embeddings

To find similar articles, let's follow a three-step plan: Get the similarity embeddings of all the article descriptions Calculate the distance between a source title and all other articles Print out the other articles closest to the source title


5. Example recommendations

Let's look for articles similar to first one, which was about Tony Blair.


Appendix: Using embeddings in more sophisticated recommenders

A more sophisticated way to build a recommender system is to train a machine learning model that takes in tens or hundreds of signals, such as item popularity or user click data. Even in this system, embeddings can be a very useful signal into the recommender, especially for items that are being 'cold started' with no user data yet (e.g., a brand new product added to the catalog without any clicks yet).


Appendix: Using embeddings to visualize similar articles

To get a sense of what our nearest neighbor recommender is doing, let's visualize the article embeddings. Although we can't plot the 2048 dimensions of each embedding vector, we can use techniques like t-SNE or PCA to compress the embeddings down into 2 or 3 dimensions, which we can chart. Before visualizing the nearest neighbors, let's visualize all of the article descriptions using t-SNE. Note that t-SNE is not deterministic, meaning that results may vary from run to run.# get embeddings for all article descriptionsembeddings = [embedding_from_string(string) for string in article_descriptions]# compress the 2048-dimensional embeddings into 2 dimensions using t-SNEtsne_components = tsne_components_from_embeddings(embeddings)# get the article labels for coloring the chartlabels = df['label'].tolist()chart_from_components( components=tsne_components, labels=labels, strings=article_descriptions, width=600, height=500, title='t-SNE components of article descriptions',)As you can see in the chart above, even the highly compressed embeddings do a good job of clustering article descriptions by category. And it's worth emphasizing: this clustering is done with no knowledge of the labels themselves! Also, if you look closely at the most egregious outliers, they are often due to mislabeling rather than poor embedding. For example, the majority of the blue World points in the green Sports cluster appear to be Sports stories.Next, let's recolor the points by whether they are a source article, its nearest neighbors, or other.# create labels for the recommended articlesdef nearest_neighbor_labels( list_of_indices: list[int], k_nearest_neighbors: int = 5) -> list[str]: '''Return a list of labels to color the k nearest neighbors.''' labels = ['Other' for _ in list_of_indices] source_index = list_of_indices[0] labels[source_index] = 'Source' for i in range(k_nearest_neighbors): nearest_neighbor_index = list_of_indices[i + 1] labels[nearest_neighbor_index] = f'Nearest neighbor (top {k_nearest_neighbors})' return labelstony_blair_labels = nearest_neighbor_labels(tony_blair_articles, k_nearest_neighbors=5)chipset_security_labels = nearest_neighbor_labels(chipset_security_articles, k_nearest_neighbors=5)# a 2D chart of nearest neighbors of the Tony Blair articlechart_from_components( components=tsne_components, labels=tony_blair_labels, strings=article_descriptions, width=600, height=500, title='Nearest neighbors of the Tony Blair article', category_orders={'label': ['Other', 'Nearest neighbor (top 5)', 'Source']},)Looking at the 2D chart above, we can see that the articles about Tony Blair are somewhat close together inside of the World news cluster. Interestingly, although the 5 nearest neighbors (red) were closest in high dimensional space, they are not the closest points in this compressed 2D space. Compressing the embeddings down to 2 dimensions discards much of their information, and the nearest neighbors in the 2D space don't seem to be as relevant as those in the full embedding space.# a 2D chart of nearest neighbors of the chipset security articlechart_from_components( components=tsne_components, labels=chipset_security_labels, strings=article_descriptions, width=600, height=500, title='Nearest neighbors of the chipset security article', category_orders={'label': ['Other', 'Nearest neighbor (top 5)', 'Source']},)For the chipset security example, the 4 closest nearest neighbors in the full embedding space remain nearest neighbors in this compressed 2D visualization. The fifth is displayed as more distant, despite being closer in the full embedding space.Should you want to, you can also make an interactive 3D plot of the embeddings with the function chart_from_components_3D. (Doing so will require recomputing the t-SNE components with n_components=3.)


Regression using the embeddings

Open in GithubRegression means predicting a number, rather than one of the categories. We will predict the score based on the embedding of the review's text. We split the dataset into a training and a testing set for all of the following tasks, so we can realistically evaluate performance on unseen data. The dataset is created in the Get_embeddings_from_dataset Notebook. We're predicting the score of the review, which is a number between 1 and 5 (1-star being negative and 5-star positive).import pandas as pdimport numpy as npfrom ast import literal_evalfrom sklearn.ensemble import RandomForestRegressorfrom sklearn.model_selection import train_test_splitfrom sklearn.metrics import mean_squared_error, mean_absolute_errordatafile_path = 'data/fine_food_reviews_with_embeddings_1k.csv'df = pd.read_csv(datafile_path)df['embedding'] = df.embedding.apply(literal_eval).apply(np.array)X_train, X_test, y_train, y_test = train_test_split(list(df.embedding.values), df.Score, test_size=0.2, random_state=42)rfr = RandomForestRegressor(n_estimators=100)rfr.fit(X_train, y_train)preds = rfr.predict(X_test)mse = mean_squared_error(y_test, preds)mae = mean_absolute_error(y_test, preds)print(f'ada-002 embedding performance on 1k Amazon reviews: mse={mse:.2f}, mae={mae:.2f}')bmse = mean_squared_error(y_test, np.repeat(y_test.mean(), len(y_test)))bmae = mean_absolute_error(y_test, np.repeat(y_test.mean(), len(y_test)))print( f'Dummy mean prediction performance on Amazon reviews: mse={bmse:.2f}, mae={bmae:.2f}')We can see that the embeddings are able to predict the scores with an average error of 0.53 per score prediction. This is roughly equivalent to predicting half of reviews perfectly, and half off by one star.You could also train a classifier to predict the label, or use the embeddings within an existing ML model to encode free text features.


Semantic text search using embeddings

Open in GithubWe can search through all our reviews semantically in a very efficient manner and at very low cost, by embedding our search query, and then finding the most similar reviews. The dataset is created in the Get_embeddings_from_dataset Notebook.import pandas as pdimport numpy as npfrom ast import literal_evaldatafile_path = 'data/fine_food_reviews_with_embeddings_1k.csv'df = pd.read_csv(datafile_path)df['embedding'] = df.embedding.apply(literal_eval).apply(np.array)Here we compare the cosine similarity of the embeddings of the query and the documents, and show top_n best matches.from openai.embeddings_utils import get_embedding, cosine_similarity# search through the reviews for a specific productdef search_reviews(df, product_description, n=3, pprint=True): product_embedding = get_embedding( product_description, engine='text-embedding-ada-002' ) df['similarity'] = df.embedding.apply(lambda x: cosine_similarity(x, product_embedding)) results = ( df.sort_values('similarity', ascending=False) .head(n) .combined.str.replace('Title: ', '') .str.replace('; Content:', ': ') ) if pprint: for r in results: print(r[:200]) print() return resultsresults = search_reviews(df, 'delicious beans', n=3)results = search_reviews(df, 'whole wheat pasta', n=3)We can search through these reviews easily. To speed up computation, we can use a special algorithm, aimed at faster search through embeddings.results = search_reviews(df, 'bad delivery', n=1)As we can see, this can immediately deliver a lot of value. In this example we show being able to quickly find the examples of delivery failures.results = search_reviews(df, 'spoilt', n=1)results = search_reviews(df, 'pet food', n=2)


User and product embeddings

Open in GithubWe calculate user and product embeddings based on the training set, and evaluate the results on the unseen test set. We will evaluate the results by plotting the user and product similarity versus the review score. The dataset is created in the Get_embeddings_from_dataset Notebook.1. Calculate user and product embeddings We calculate these embeddings simply by averaging all the reviews about the same product or written by the same user within the training set.import pandas as pdimport numpy as npfrom sklearn.model_selection import train_test_splitfrom ast import literal_evaldf = pd.read_csv('output/embedded_babbage_similarity_50k.csv', index_col=0) # note that you will need to generate this file to run the code belowdf['babbage_similarity'] = df.babbage_similarity.apply(literal_eval).apply(np.array)X_train, X_test, y_train, y_test = train_test_split(df, df.Score, test_size = 0.2, random_state=42)user_embeddings = X_train.groupby('UserId').babbage_similarity.apply(np.mean)prod_embeddings = X_train.groupby('ProductId').babbage_similarity.apply(np.mean)len(user_embeddings), len(prod_embeddings)We can see that most of the users and products appear within the 50k examples only once.2. Evaluate the embeddings To evaluate the recommendations, we look at the similarity of the user and product embeddings amongst the reviews in the unseen test set. We calculate the cosine distance between the user and product embeddings, which gives us a similarity score between 0 and 1. We then normalize the scores to be evenly split between 0 and 1, by calculating the percentile of the similarity score amongst all predicted scores.from openai.embeddings_utils import cosine_similarity# evaluate embeddings as recommendations on X_testdef evaluate_single_match(row): user_id = row.UserId product_id = row.ProductId try: user_embedding = user_embeddings[user_id] product_embedding = prod_embeddings[product_id] similarity = cosine_similarity(user_embedding, product_embedding) return similarity except Exception as e: return np.nanX_test['cosine_similarity'] = X_test.apply(evaluate_single_match, axis=1)X_test['percentile_cosine_similarity'] = X_test.cosine_similarity.rank(pct=True)2.1 Visualize cosine similarity by review score We group the cosine similarity scores by the review score, and plot the distribution of cosine similarity scores for each review score.import matplotlib.pyplot as pltimport statsmodels.api as smcorrelation = X_test[['percentile_cosine_similarity', 'Score']].corr().values[0,1]print('Correlation between user & vector similarity percentile metric and review number of stars (score): %.2f%%' % (100*correlation))# boxplot of cosine similarity for each scoreX_test.boxplot(column='percentile_cosine_similarity', by='Score')plt.title('')plt.show()plt.close()We can observe a weak trend, showing that the higher the similarity score between the user and the product embedding, the higher the review score. Therefore, the user and product embeddings can weakly predict the review score - even before the user receives the product! Because this signal works in a different way than the more commonly used collaborative filtering, it can act as an additional feature to slightly improve the performance on existing problems.


1. Calculate user and product embeddings

We calculate these embeddings simply by averaging all the reviews about the same product or written by the same user within the training set.


2.1 Visualize cosine similarity by review score

We group the cosine similarity scores by the review score, and plot the distribution of cosine similarity scores for each review score.


2. Evaluate the embeddings

To evaluate the recommendations, we look at the similarity of the user and product embeddings amongst the reviews in the unseen test set. We calculate the cosine distance between the user and product embeddings, which gives us a similarity score between 0 and 1. We then normalize the scores to be evenly split between 0 and 1, by calculating the percentile of the similarity score amongst all predicted scores.


Visualizing the embeddings in 2D

Open in GithubWe will use t-SNE to reduce the dimensionality of the embeddings from 1536 to 2. Once the embeddings are reduced to two dimensions, we can plot them in a 2D scatter plot. The dataset is created in the Get_embeddings_from_dataset Notebook.1. Reduce dimensionality We reduce the dimensionality to 2 dimensions using t-SNE decomposition.import pandas as pdfrom sklearn.manifold import TSNEimport numpy as npfrom ast import literal_eval# Load the embeddingsdatafile_path = 'data/fine_food_reviews_with_embeddings_1k.csv'df = pd.read_csv(datafile_path)# Convert to a list of lists of floatsmatrix = np.array(df.embedding.apply(literal_eval).to_list())# Create a t-SNE model and transform the datatsne = TSNE(n_components=2, perplexity=15, random_state=42, init='random', learning_rate=200)vis_dims = tsne.fit_transform(matrix)vis_dims.shape2. Plotting the embeddings We colour each review by its star rating, ranging from red to green.We can observe a decent data separation even in the reduced 2 dimensions.import matplotlib.pyplot as pltimport matplotlibimport numpy as npcolors = ['red', 'darkorange', 'gold', 'turquoise', 'darkgreen']x = [x for x,y in vis_dims]y = [y for x,y in vis_dims]color_indices = df.Score.values - 1colormap = matplotlib.colors.ListedColormap(colors)plt.scatter(x, y, c=color_indices, cmap=colormap, alpha=0.3)for score in [0,1,2,3,4]: avg_x = np.array(x)[df.Score-1==score].mean() avg_y = np.array(y)[df.Score-1==score].mean() color = colors[score] plt.scatter(avg_x, avg_y, marker='x', color=color, s=100)plt.title('Amazon ratings visualized in language using t-SNE')


2. Plotting the embeddings

We colour each review by its star rating, ranging from red to green.


1. Reduce dimensionality

We reduce the dimensionality to 2 dimensions using t-SNE decomposition.


Visualizing embeddings in 3D

Open in GithubThe example uses PCA to reduce the dimensionality fo the embeddings from 1536 to 3. Then we can visualize the data points in a 3D plot. The small dataset dbpedia_samples.jsonl is curated by randomly sampling 200 samples from DBpedia validation dataset.


1. Load the dataset and query embeddings

import pandas as pdsamples = pd.read_json('data/dbpedia_samples.jsonl', lines=True)categories = sorted(samples['category'].unique())print('Categories of DBpedia samples:', samples['category'].value_counts())samples.head() .dataframe tbody tr th:only-of-type { vertical-align: middle; } .dataframe tbody tr th { vertical-align: top; } .dataframe thead th { text-align: right; } text category 0 Morada Limited is a textile company based in ... Company 1 The Armenian Mirror-Spectator is a newspaper ... WrittenWork 2 Mt. Kinka (金華山 Kinka-zan) also known as Kinka... NaturalPlace 3 Planning the Play of a Bridge Hand is a book ... WrittenWork 4 Wang Yuanping (born 8 December 1976) is a ret... Athlete from openai.embeddings_utils import get_embeddings# NOTE: The following code will send a query of batch size 200 to /embeddingsmatrix = get_embeddings(samples['text'].to_list(), engine='text-embedding-ada-002')


3. Plot the embeddings of lower dimensionality

%matplotlib widgetimport matplotlib.pyplot as pltimport numpy as npfig = plt.figure(figsize=(10, 5))ax = fig.add_subplot(projection='3d')cmap = plt.get_cmap('tab20')# Plot each sample category individually such that we can set label name.for i, cat in enumerate(categories): sub_matrix = np.array(samples[samples['category'] == cat]['embed_vis'].to_list()) x=sub_matrix[:, 0] y=sub_matrix[:, 1] z=sub_matrix[:, 2] colors = [cmap(i/len(categories))] * len(sub_matrix) ax.scatter(x, y, zs=z, zdir='z', c=colors, label=cat)ax.set_xlabel('x')ax.set_ylabel('y')ax.set_zlabel('z')ax.legend(bbox_to_anchor=(1.1, 1))


2. Reduce the embedding dimensionality

from sklearn.decomposition import PCApca = PCA(n_components=3)vis_dims = pca.fit_transform(matrix)samples['embed_vis'] = vis_dims.tolist()


Zero-shot classification with embeddings

Open in GithubIn this notebook we will classify the sentiment of reviews using embeddings and zero labeled data! The dataset is created in the Get_embeddings_from_dataset Notebook. We'll define positive sentiment to be 4- and 5-star reviews, and negative sentiment to be 1- and 2-star reviews. 3-star reviews are considered neutral and we won't use them for this example. We will perform zero-shot classification by embedding descriptions of each class and then comparing new samples to those class embeddings.# importsimport pandas as pdimport numpy as npfrom ast import literal_evalfrom sklearn.metrics import classification_report# parametersEMBEDDING_MODEL = 'text-embedding-ada-002'# load datadatafile_path = 'data/fine_food_reviews_with_embeddings_1k.csv'df = pd.read_csv(datafile_path)df['embedding'] = df.embedding.apply(literal_eval).apply(np.array)# convert 5-star rating to binary sentimentdf = df[df.Score != 3]df['sentiment'] = df.Score.replace({1: 'negative', 2: 'negative', 4: 'positive', 5: 'positive'})Zero-Shot Classification To perform zero shot classification, we want to predict labels for our samples without any training. To do this, we can simply embed short descriptions of each label, such as positive and negative, and then compare the cosine distance between embeddings of samples and label descriptions. The highest similarity label to the sample input is the predicted label. We can also define a prediction score to be the difference between the cosine distance to the positive and to the negative label. This score can be used for plotting a precision-recall curve, which can be used to select a different tradeoff between precision and recall, by selecting a different threshold.from openai.embeddings_utils import cosine_similarity, get_embeddingfrom sklearn.metrics import PrecisionRecallDisplaydef evaluate_embeddings_approach( labels = ['negative', 'positive'], model = EMBEDDING_MODEL,): label_embeddings = [get_embedding(label, engine=model) for label in labels] def label_score(review_embedding, label_embeddings): return cosine_similarity(review_embedding, label_embeddings[1]) - cosine_similarity(review_embedding, label_embeddings[0]) probas = df['embedding'].apply(lambda x: label_score(x, label_embeddings)) preds = probas.apply(lambda x: 'positive' if x>0 else 'negative') report = classification_report(df.sentiment, preds) print(report) display = PrecisionRecallDisplay.from_predictions(df.sentiment, probas, pos_label='positive') _ = display.ax_.set_title('2-class Precision-Recall curve')evaluate_embeddings_approach(labels=['negative', 'positive'], model=EMBEDDING_MODEL)We can see that this classifier already performs extremely well. We used similarity embeddings, and the simplest possible label name. Let's try to improve on this by using more descriptive label names, and search embeddings.evaluate_embeddings_approach(labels=['An Amazon review with a negative sentiment.', 'An Amazon review with a positive sentiment.'])Using the search embeddings and descriptive names leads to an additional improvement in performance.evaluate_embeddings_approach(labels=['An Amazon review with a negative sentiment.', 'An Amazon review with a positive sentiment.'])As shown above, zero-shot classification with embeddings can lead to great results, especially when the labels are more descriptive than just simple words.


Zero-Shot Classification

To perform zero shot classification, we want to predict labels for our samples without any training. To do this, we can simply embed short descriptions of each label, such as positive and negative, and then compare the cosine distance between embeddings of samples and label descriptions.The highest similarity label to the sample input is the predicted label. We can also define a prediction score to be the difference between the cosine distance to the positive and to the negative label. This score can be used for plotting a precision-recall curve, which can be used to select a different tradeoff between precision and recall, by selecting a different threshold.


Translate a book writen in LaTeX from Slovenian into English

Open in GithubWith permission of the author, we will demonstrate how to translate the book Euclidean Plane Geometry, written by Milan Mitrović from Slovenian into English, without modifying any of the LaTeX commands. To achieve this, we will first split the book into chunks, each roughly a page long, then translate each chunk into English, and finally stitch them back together.


1. Read in the data

import openaifrom transformers import GPT2Tokenizer# OpenAI GPT-2 tokenizer is the same as GPT-3 tokenizer# we use it to count the number of tokens in the texttokenizer = GPT2Tokenizer.from_pretrained('gpt2')with open('data/geometry_slovenian.tex', 'r') as f: text = f.read()


1.1 Count the tokens in each chunk

chunks = text.split('\n\n')ntokens = []for chunk in chunks: ntokens.append(len(tokenizer.encode(chunk)))max(ntokens)It turns out that a double newline is a good separator in this case, in order not to break the flow of the text. Also no individual chunk is larger than 1500 tokens. The model we will use is text-davinci-002, which has a limit of 4096 tokens, so we don't need to worry about breaking the chunks down further. We will group the shorter chunks into chunks of around 1000 tokens, to increase the coherence of the text, and decrease the frequency of breaks within the text.def group_chunks(chunks, ntokens, max_len=1000, hard_max_len=3000): ''' Group very short chunks, to form approximately page long chunks. ''' batches = [] cur_batch = '' cur_tokens = 0 # iterate over chunks, and group the short ones together for chunk, ntoken in zip(chunks, ntokens): # discard chunks that exceed hard max length if ntoken > hard_max_len: print(f'Warning: Chunk discarded for being too long ({ntoken} tokens > {hard_max_len} token limit). Preview: '{chunk[:50]}...'') continue # if room in current batch, add new chunk if cur_tokens + 1 + ntoken <= max_len: cur_batch += '\n\n' + chunk cur_tokens += 1 + ntoken # adds 1 token for the two newlines # otherwise, record the batch and start a new one else: batches.append(cur_batch) cur_batch = chunk cur_tokens = ntoken if cur_batch: # add the last batch if it's not empty batches.append(cur_batch) return batcheschunks = group_chunks(chunks, ntokens)len(chunks)Notice that adding a sample untranslated and translated first command, where only the content of the chapter name needs to be translated, helps to get more consistent results. The format of the prompt sent to the model consists of: A high level instruction to translate only the text, but not commands into the desired language A sample untranslated command, where only the content of the chapter name needs to be translated The chunk of text to be translated The translated sample command from 2, which shows the model the beginning of the translation process The expected output is the translated chunk of text.def translate_chunk(chunk, engine='text-davinci-002', dest_language='English', sample_translation=('\poglavje{Osnove Geometrije} \label{osn9Geom}', '\poglavje{The basics of Geometry} \label{osn9Geom}') ): prompt = f'''Translate only the text from the following LaTeX document into {dest_language}. Leave all LaTeX commands unchanged '''{sample_translation[0]}{chunk}'''{sample_translation[1]}''' response = openai.Completion.create( prompt=prompt, engine=engine, temperature=0, top_p=1, max_tokens=1500, ) result = response['choices'][0]['text'].strip() result = result.replace(''''', '') # remove the double quotes, as we used them to surround the text return resultprint(translate_chunk(chunks[800], engine='text-davinci-002', dest_language='English'))We can see here that this one chunk in particular translates only the text, but leaves LaTeX commands intact. Let's now translate all the chunks in the book - this will take 2-3 hours, as we're processing requests sequentially.dest_language = 'English'translated_chunks = []for i, chunk in enumerate(chunks): print(str(i+1) + ' / ' + str(len(chunks))) # translate each chunk translated_chunks.append(translate_chunk(chunk, engine='text-davinci-002', dest_language=dest_language))# join the chunks togetherresult = '\n\n'.join(translated_chunks)# save the final resultwith open(f'data/geometry_{dest_language}.tex', 'w') as f: f.write(result)


Fine-Tuned Q&A - Collect Data

Open in GithubNote: To answer questions based on text documents, we recommend the procedure in Question Answering using Embeddings.  Some of the code below may rely on deprecated API endpoints.1.  Collect Wikipedia data about Olympic Games 2020 The idea of this project is to create a question answering model, based on a few paragraphs of provided text.  Base GPT-3 models do a good job at answering questions when the answer is contained within the paragraph, however if the answer isn't contained, the base models tend to try their best to answer anyway, often leading to confabulated answers.  To create a model which answers questions only if there is sufficient context for doing so, we first create a dataset of questions and answers based on paragraphs of text.  In order to train the model to answer only when the answer is present, we also add adversarial examples, where the question doesn't match the context.  In those cases, we ask the model to output 'No sufficient context for answering the question'.  We will perform this task in three notebooks: The first (this) notebook focuses on collecting recent data, which GPT-3 didn't see during it's pre-training.  We picked the topic of Olympic Games 2020 (which actually took place in the summer of 2021), and downloaded 713 unique pages.  We organized the dataset by individual sections, which will serve as context for asking and answering the questions.  The second notebook will utilize Davinci-instruct to ask a few questions based on a Wikipedia section, as well as answer those questions, based on that section.  The third notebook will utilize the dataset of context, question and answer pairs to additionally create adversarial questions and context pairs, where the question was not generated on that context.  In those cases the model will be prompted to answer 'No sufficient context for answering the question'.  We will also train a discriminator model, which predicts whether the question can be answered based on the context or not


Fine-Tuned Q&A - Collect Data

1.1 Data extraction using the wikipedia API Extracting the data will take about half an hour, and processing will likely take about as much.import pandas as pdimport wikipediadef filter_olympic_2020_titles(titles): ''' Get the titles which are related to Olympic games hosted in 2020, given a list of titles ''' titles = [title for title in titles if '2020' in title and 'olympi' in title.lower()] return titlesdef get_wiki_page(title): ''' Get the wikipedia page given a title ''' try: return wikipedia.page(title) except wikipedia.exceptions.DisambiguationError as e: return wikipedia.page(e.options[0]) except wikipedia.exceptions.PageError as e: return Nonedef recursively_find_all_pages(titles, titles_so_far=set()): ''' Recursively find all the pages that are linked to the Wikipedia titles in the list ''' all_pages = [] titles = list(set(titles) - titles_so_far) titles = filter_olympic_2020_titles(titles) titles_so_far.update(titles) for title in titles: page = get_wiki_page(title) if page is None: continue all_pages.append(page) new_pages = recursively_find_all_pages(page.links, titles_so_far) for pg in new_pages: if pg.title not in [p.title for p in all_pages]: all_pages.append(pg) titles_so_far.update(page.links) return all_pagespages = recursively_find_all_pages(['2020 Summer Olympics'])len(pages)1.2 Filtering the Wikipedia pages and splitting them into sections by headings We remove sections unlikely to contain textual information, and ensure that each section is not longer than the token limitimport refrom typing import Setfrom transformers import GPT2TokenizerFastimport numpy as npfrom nltk.tokenize import sent_tokenizetokenizer = GPT2TokenizerFast.from_pretrained('gpt2')def count_tokens(text: str) -> int: '''count the number of tokens in a string''' return len(tokenizer.encode(text))def reduce_long( long_text: str, long_text_tokens: bool = False, max_len: int = 590) -> str: ''' Reduce a long text to a maximum of `max_len` tokens by potentially cutting at a sentence end ''' if not long_text_tokens: long_text_tokens = count_tokens(long_text) if long_text_tokens > max_len: sentences = sent_tokenize(long_text.replace('\n', ' ')) ntokens = 0 for i, sentence in enumerate(sentences): ntokens += 1 + count_tokens(sentence) if ntokens > max_len: return '


Fine-Tuned Q&A - Collect Data

'.join(sentences[:i]) + '.' return long_textdiscard_categories = ['See also', 'References', 'External links', 'Further reading', 'Footnotes', 'Bibliography', 'Sources', 'Citations', 'Literature', 'Footnotes', 'Notes and references', 'Photo gallery', 'Works cited', 'Photos', 'Gallery', 'Notes', 'References and sources', 'References and notes',]def extract_sections( wiki_text: str, title: str, max_len: int = 1500, discard_categories: Set[str] = discard_categories,) -> str: ''' Extract the sections of a Wikipedia page, discarding the references and other low information sections ''' if len(wiki_text) == 0: return [] # find all headings and the corresponding contents headings = re.findall('==+ .* ==+', wiki_text) for heading in headings: wiki_text = wiki_text.replace(heading, '==+ !! ==+') contents = wiki_text.split('==+ !! ==+') contents = [c.strip() for c in contents] assert len(headings) == len(contents) - 1 cont = contents.pop(0).strip() outputs = [(title, 'Summary', cont, count_tokens(cont)+4)] # discard the discard categories, accounting for a tree structure max_level = 100 keep_group_level = max_level remove_group_level = max_level nheadings, ncontents = [], [] for heading, content in zip(headings, contents): plain_heading = ' '.join(heading.split(' ')[1:-1]) num_equals = len(heading.split(' ')[0]) if num_equals <= keep_group_level: keep_group_level = max_level if num_equals > remove_group_level: if ( num_equals <= keep_group_level ): continue keep_group_level = max_level if plain_heading in discard_categories: remove_group_level = num_equals keep_group_level = max_level continue nheadings.append(heading.replace('=', '').strip()) ncontents.append(content) remove_group_level = max_level # count the tokens of each section ncontent_ntokens = [ count_tokens(c) + 3 + count_tokens(' '.join(h.split(' ')[1:-1])) - (1 if len(c) == 0 else 0) for h, c in zip(nheadings, ncontents) ] # Create a tuple of (title, section_name, content, number of tokens) outputs += [(title, h, c, t) if t<max_len else (title, h, reduce_long(c, max_len), count_tokens(reduce_long(c,max_len))) for h, c, t in zip(nheadings, ncontents, ncontent_ntokens)] return outputs# Example page being processed into sectionsbermuda_page = get_wiki_page('Bermuda at the 2020 Summer Olympics')ber = extract_sections(bermuda_page.content, bermuda_page.title)# Example sectionber[-1]


1. Collect Wikipedia data about Olympic Games 2020

The idea of this project is to create a question answering model, based on a few paragraphs of provided text. Base GPT-3 models do a good job at answering questions when the answer is contained within the paragraph, however if the answer isn't contained, the base models tend to try their best to answer anyway, often leading to confabulated answers.To create a model which answers questions only if there is sufficient context for doing so, we first create a dataset of questions and answers based on paragraphs of text. In order to train the model to answer only when the answer is present, we also add adversarial examples, where the question doesn't match the context. In those cases, we ask the model to output 'No sufficient context for answering the question'.We will perform this task in three notebooks: The first (this) notebook focuses on collecting recent data, which GPT-3 didn't see during it's pre-training. We picked the topic of Olympic Games 2020 (which actually took place in the summer of 2021), and downloaded 713 unique pages. We organized the dataset by individual sections, which will serve as context for asking and answering the questions. The second notebook will utilize Davinci-instruct to ask a few questions based on a Wikipedia section, as well as answer those questions, based on that section. The third notebook will utilize the dataset of context, question and answer pairs to additionally create adversarial questions and context pairs, where the question was not generated on that context. In those cases the model will be prompted to answer 'No sufficient context for answering the question'. We will also train a discriminator model, which predicts whether the question can be answered based on the context or not.


1.2 Filtering the Wikipedia pages and splitting them into sections by headings

We remove sections unlikely to contain textual information, and ensure that each section is not longer than the token limit


1.2.1 We create a dataset and filter out any sections with fewer than 40 tokens, as those are unlikely to contain enough context to ask a good question.

res = []for page in pages: res += extract_sections(page.content, page.title)df = pd.DataFrame(res, columns=['title', 'heading', 'content', 'tokens'])df = df[df.tokens>40]df = df.drop_duplicates(['title','heading'])df = df.reset_index().drop('index',axis=1) # reset indexdf.head() .dataframe tbody tr th:only-of-type { vertical-align: middle; } .dataframe tbody tr th { vertical-align: top; } .dataframe thead th { text-align: right; } title heading content tokens 0 2020 Summer Olympics Summary The 2020 Summer Olympics (Japanese: 2020年夏季オリン... 713 1 2020 Summer Olympics Host city selection The International Olympic Committee (IOC) vote... 126 2 2020 Summer Olympics Impact of the COVID-19 pandemic In January 2020, concerns were raised about th... 369 3 2020 Summer Olympics Qualifying event cancellation and postponement Concerns about the pandemic began to affect qu... 298 4 2020 Summer Olympics Effect on doping tests Mandatory doping tests were being severely res... 163 Save the section dataset We will save the section dataset, for the next notebookdf.to_csv('olympics-data/olympics_sections.csv', index=False)


1.1 Data extraction using the wikipedia API

Extracting the data will take about half an hour, and processing will likely take about as much.


Save the section dataset

We will save the section dataset, for the next notebook


1.3 (Optional) Exploring the data

df.title.value_counts().head()There appear to be winter and summer Olympics 2020. We chose to leave a little ambiguity and noise in the dataset, even though we were interested in only Summer Olympics 2020.df.title.str.contains('Summer').value_counts()df.title.str.contains('Winter').value_counts()import pandas as pdfrom matplotlib import pyplot as pltdf = pd.read_csv('olympics-data/olympics_sections.csv')df[['tokens']].hist()# add axis descriptions and titleplt.xlabel('Number of tokens')plt.ylabel('Number of Wikipedia sections')plt.title('Distribution of number of tokens in Wikipedia sections')plt.show()We can see that the majority of section are fairly short (less than 500 tokens).


Fine-Tuned Q&A - Create Q&A

WARNING: This step will last a long time, and consume a lot of tokens, as it calls davinci-instruct for every section to answer all the questions.def get_answers(row): try: response = openai.Completion.create( engine='davinci-instruct-beta-v3', prompt=f'Write answer based on the text below\n\nText: {row.context}\n\nQuestions:\n{row.questions}\n\nAnswers:\n1.', temperature=0, max_tokens=257, top_p=1, frequency_penalty=0, presence_penalty=0 ) return response['choices'][0]['text'] except Exception as e: print (e) return ''df['answers']= df.apply(get_answers, axis=1)df['answers'] = '1.' + df.answersdf = df.dropna().reset_index().drop('index',axis=1)print(df[['answers']].values[0][0])These are the answers to the questions above based on the context around the host city selection .  We can see that answers 3-5 contain the correct answer, but instead of answering the question directly, the answer is a verbatim extraction.  Despite these occasional lower quality answers, we will show that the model can learn the task reasonably well, given a high number of examples.2.4 Save the Olympics Q&A dataset based on Wikipedia sections We save the file for use in the next notebookdf.to_csv('olympics-data/olympics_qa.csv', index=False)2.5 Search file (DEPRECATED) We create a search file (API reference), which can be used to retrieve the relevant context when a question is asked.  DEPRECATED: The /search endpoint is deprecated in favour of using embeddings.  Embeddings are cheaper, faster and can support a better search experience.  See Question Answering Guide for a search implementation using the embeddingsdf = df[df.tokens<2000]df[['context', 'tokens']].rename(columns={'context':'text','tokens':'metadata'}).to_json('olympics-data/olympics_search.jsonl', orient='records', lines=True)search_file = openai.File.create( file=open('olympics-data/olympics_search.jsonl'), purpose='search')olympics_search_fileid = search_file['id']2.6 Answer questions based on the context provided We will use a simple implementation of the answers endpoint.  This works by simply using the /search endpoint, which searches over an indexed file to obtain the relevant sections which can be included in the context, following by a question and answering prompt given a specified model.from answers_with_ft import create_context, answer_questionprint(create_context('Where did women's 4 x 100 metres relay event take place during the 2020 Summer Olympics?', olympics_search_fileid, max_len=400))answer_question(olympics_search_fileid, 'davinci-instruct-beta-v3', 'Where did women's 4 x 100 metres relay event take place during the 2020 Summer Olympics?')After we fine-tune the model for Q&A we'll be able to use it instead of davinci-instruct-beta-v3, to obtain better answers when the question can't be answered based on the context.  We see a downside of davinci-instruct-beta-v3, which always attempts to answer the question, regardless of the relevant context being present or not.  (Note the second question is asking about a future event, set in 2024.)answer_question(olympics_search_fileid, 'davinci-instruct-beta-v3', 'Where did women's 4 x 100 metres relay event take place during the 2048 Summer Olympics?', max_len=1000)We can see that davinci has a tendency to answer the question, even if the question can't be answered given the context provided.  Note the question asked regarding 2048 Summer Olympics, which didn't happen yet, and the retrieved content has only returned results for 2020.


Fine-Tuned Q&A - Create Q&A

Open in GithubNote: To answer questions based on text documents, we recommend the procedure in Question Answering using Embeddings.  Some of the code below may rely on deprecated API endpoints.2.  Creating a synthetic Q&A dataset We use davinci-instruct-beta-v3, a model specialized in following instructions, to create questions based on the given context.  Then we also use davinci-instruct-beta-v3 to answer those questions, given the same context.  This is expensive, and will also take a long time, as we call the davinci engine for each section.  You can simply download the final dataset instead.  We're using the dataset created using the previous notebook2.1 Read in the data, and create a context Create a context by concatenating the title, the heading and the content of that sectionimport pandas as pddf = pd.read_csv('olympics-data/olympics_sections.csv')df['context'] = df.title + '\n' + df.heading + '\n\n' + df.contentdf.head() .dataframe tbody tr th:only-of-type { vertical-align: middle; } .dataframe tbody tr th { vertical-align: top; } .dataframe thead th { text-align: right; } title heading content tokens context 0 2020 Summer Olympics Summary The 2020 Summer Olympics (Japanese: 2020年夏季オリン...  713 2020 Summer Olympics\nSummary\n\nThe 2020 Summ...  1 2020 Summer Olympics Host city selection The International Olympic Committee (IOC) vote...  126 2020 Summer Olympics\nHost city selection\n\nT...  2 2020 Summer Olympics Impact of the COVID-19 pandemic In January 2020, concerns were raised about th...  369 2020 Summer Olympics\nImpact of the COVID-19 p...  3 2020 Summer Olympics Qualifying event cancellation and postponement Concerns about the pandemic began to affect qu...  298 2020 Summer Olympics\nQualifying event cancell...  4 2020 Summer Olympics Effect on doping tests Mandatory doping tests were being severely res...  163 2020 Summer Olympics\nEffect on doping tests\n...  2.2 Create questions based on the context Use davinci-instruct to generate a number of plausible questions relating to the Wikipedia section contents.  Note: We have used temperature=0, but it may be beneficial to experiment with a higher temperature to get a higher diversity of questions.  WARNING: This step will last a long time, and consume a lot of tokens, as it calls davinci-instruct for every section to generate a number of questions.import openaidef get_questions(context): try: response = openai.Completion.create( engine='davinci-instruct-beta-v3', prompt=f'Write questions based on the text below\n\nText: {context}\n\nQuestions:\n1.', temperature=0, max_tokens=257, top_p=1, frequency_penalty=0, presence_penalty=0, stop=['\n\n'] ) return response['choices'][0]['text'] except: return ''df['questions']= df.context.apply(get_questions)df['questions'] = '1.' + df.questionsprint(df[['questions']].values[0][0])The prompt is designed to generate a number of questions.  Example questions above were generated based on the summary section of the 2020 Summer Olympics page.  We can observe that the questions 3 and 5 above repeat.  Sometimes the generated questions could be ambiguous without the context.  We will show that even despite these limitations we can create a successful model.print(df.content.values[0])2.3 Create answers based on the context Use davinci-instruct to answer the questions given the relevant Wikipedia section contents Note: We have used temperature=0, but it may be beneficial to experiment with a higher temperature to get a higher diversity of questions


2.2 Create questions based on the context

Use davinci-instruct to generate a number of plausible questions relating to the Wikipedia section contents.Note: We have used temperature=0, but it may be beneficial to experiment with a higher temperature to get a higher diversity of questions.WARNING: This step will last a long time, and consume a lot of tokens, as it calls davinci-instruct for every section to generate a number of questions.


2.1 Read in the data, and create a context

Create a context by concatenating the title, the heading and the content of that section


2.3 Create answers based on the context

Use davinci-instruct to answer the questions given the relevant Wikipedia section contentsNote: We have used temperature=0, but it may be beneficial to experiment with a higher temperature to get a higher diversity of questions.WARNING: This step will last a long time, and consume a lot of tokens, as it calls davinci-instruct for every section to answer all the questions.


2. Creating a synthetic Q&A dataset

We use davinci-instruct-beta-v3, a model specialized in following instructions, to create questions based on the given context. Then we also use davinci-instruct-beta-v3 to answer those questions, given the same context.This is expensive, and will also take a long time, as we call the davinci engine for each section. You can simply download the final dataset instead.We're using the dataset created using the previous notebook


2.4 Save the Olympics Q&A dataset based on Wikipedia sections

We save the file for use in the next notebook


2.7 (Optional) Investigation into how likely the search endpoint is to return the relevant context

def check_context(title, heading, question, max_len=1800, search_model='ada', max_rerank=10): ''' Evaluate the performance of the search model in retrieving the correct context Parameters ---------- title: str The title of the Wikipedia page heading: str The heading of the Wikipedia section qusetion: str The question max_len: int The maximum length of the context search_model: str The search model to use - `ada` is most cost effective max_rerank: int The maximum number of reranking documents to use the search model on Returns ------- rank: int The rank of the correct context token_length: int The number of tokens needed to obtain the correct context ''' try: results = openai.Engine(search_model).search( search_model=search_model, query=question, max_rerank=max_rerank, file=olympics_search_fileid, return_metadata=True ) index=-1 returns = [] cur_len = 0 for result in results['data']: cur_len += int(result['metadata']) + 4 # we add 4 tokens for the separator `\n\n###\n\n` if cur_len > max_len: break returns.append(result['text']) res = result['text'].split('\n') if res[0] == title and res[1] == heading: index = len(returns) - 1 break return index, cur_len except Exception as e: #print (e) return []print(check_context('Athletics at the 2020 Summer Olympics – Women's 4 × 100 metres relay', 'Summary', 'Where did women's 4 x 100 metres relay event take place during the 2020 Summer Olympics?', max_len=10000))We utilize the generated questions based on context to estimate how often we can retrieve the original context.  These questions are noisy, so this is not a perfect estimate.  Our questions and answers are prefixed with numbered bullet points, however due to the way they were generated, they are missing the first number, hence we add '1.' to the list of questions (and answers).  We calculate the rank of the section retrieved using ada search, and the number of tokens in the context needed to retrieve the relevant section in full.ada_results = df.apply(lambda x: [ check_context( x.title, x.heading, q[3:], # remove the number prefix max_len=1000000, # set a large number to get the full context search_model='ada', max_rerank=200, ) for q in (x.questions).split('\n') # split the questions if len(q) >10 # remove the empty questions ], axis=1)ada_results.head()out = pd.concat([ada_results], axis=1)out.columns = ['ada']out.to_csv('olympics-data/search_engine_results.csv')def expand_lists(out): ''' Expand a pandas series containing lists into a series, where each list element becomes a value on its own Input is a row per paragraph, which has multiple questions Output is a row per question ''' cols = [pd.DataFrame(out[name].tolist()).stack().reset_index(level=1, drop=True).rename(name) for name in out.columns] return pd.concat(cols, axis=1)out_expanded = expand_lists(out)out_expanded['rank'] = out_expanded.ada.apply(lambda x: x[0] if x != [] else -2)out_expanded['tokens'] = out_expanded.ada.apply(lambda x: x[1] if x != [] else -2)within_2k = (out_expanded.tokens < 2000).mean()print(f'{within_2k*100:.1f}% of relevant paragraphs are retrieved within the first 2k tokens')The relevant context can be obtained 74% of the time on this datasetoutside_200 = (out_expanded['rank'] == -1).mean()print(f'{outside_200*100:.1f}% of relevant paragraphs are not retrieved within the first 200 results')7.4% of the time, this is due to the keyword search part of the search algorithm not retrieving the relevant context within the first 200 results


18.3% of the time this is due to the semantic search not placing the relevant context within the first 2000 tokens.import matplotlib.pyplot as plt# plot a histogram, and add axis descriptions and titleout_expanded[(out_expanded['rank'] >=0)&(out_expanded['rank'] <30)]['rank'].hist(bins=29)plt.xlabel('rank')plt.ylabel('count')plt.title('Histogram of ranks of retrieved paragraphs')plt.show()out_expanded[(out_expanded.tokens>=0)&(out_expanded.tokens < 2000)]['tokens'].hist(bins=29)plt.xlabel('tokens')plt.ylabel('count')plt.title('Histogram of the number of minimum tokens needed')plt.show()We can observe that the context is most likely to be returned as one of the first results, and most likely to be returned within the first 200-500 tokens.# normalized value_countsout_expanded['rank'].value_counts(normalize=True).sort_index()[:13]probabilities of the relevant context being returned at each rank .  (-2 means a processing error, -1 means the rank is >200)


2.6 Answer questions based on the context provided

We will use a simple implementation of the answers endpoint. This works by simply using the /search endpoint, which searches over an indexed file to obtain the relevant sections which can be included in the context, following by a question and answering prompt given a specified model.


2.5 Search file (DEPRECATED)

We create a search file (API reference), which can be used to retrieve the relevant context when a question is asked.DEPRECATED: The /search endpoint is deprecated in favour of using embeddings. Embeddings are cheaper, faster and can support a better search experience. See Question Answering Guide for a search implementation using the embeddings