Using embeddings to fuzzy match databases

linking_with_embeddings.py

import linktransformer as lt
import pandas as pd
import os

# Read in the two datasets
df1 = pd.read_csv("test1.csv")
df2 = pd.read_csv("test2.csv")

df_lm_matched = lt.merge(df1, df2, merge_type='1:m', model="all-MiniLM-L6-v2", left_on="HS", right_on="name")
print(df_lm_matched.head(10))

linking_with_embeddings.pyout

   id                            HS
0   1        Timpanogos High School
1   2                 Timpanogos HS
2   3                    Timpanogos
3   4              Archbishop Mitty
4   5    Archbishop McCarthy (Fla.)
5   6            Archbishop Carroll
6   7          Bishop Ryan Catholic
7   8               Bishop Shanahan
8   9  Bishop Hendricken (American)
9  10                 Bishop Feehan
    schid                                           name
0  100000                                TIMPANOGOS HIGH
1  111111                                  TIMPVIEW HIGH
2  121212                                    TIMPSON H S
3  123456                   ARCHBISHOP HANNA HIGH SCHOOL
4  135792                   ARCHBISHOP MITTY HIGH SCHOOL
5  151515                 ARCHBISHOP RIORDAN HIGH SCHOOL
6  161616                     BISHOP ALEMANY HIGH SCHOOL
7  171717               BISHOP AMAT MEMORIAL HIGH SCHOOL
8  181818  BISHOP CONATY-OUR LADY OF LORETTO HIGH SCHOOL
9  191919                BISHOP GARCIA DIEGO HIGH SCHOOL
2023-09-16 17:09:45 - Load pretrained SentenceTransformer: all-MiniLM-L6-v2
2023-09-16 17:09:46 - Use pytorch device: cpu
Batches: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 31.92it/s]
Batches: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00, 14.00it/s] 
LM matched on key columns - left: HS_x, right: name_y
   id                            HS  id_lt_x   schid                                 name  id_lt_y     score
0   1        Timpanogos High School        0  100000                      TIMPANOGOS HIGH        0  0.756922
1   2                 Timpanogos HS        1  100000                      TIMPANOGOS HIGH        0  0.864616
2   3                    Timpanogos        2  100000                      TIMPANOGOS HIGH        0  0.874583
3   4              Archbishop Mitty        3  135792         ARCHBISHOP MITTY HIGH SCHOOL        4  0.787486
4   5    Archbishop McCarthy (Fla.)        4  252525      ARCHBISHOP MCCARTHY HIGH SCHOOL       14  0.689855
5   6            Archbishop Carroll        5  575757  ARCHBISHOP JOHN CARROLL HIGH SCHOOL      117  0.754913
6   7          Bishop Ryan Catholic        6  523459          BISHOP RYAN CATHOLIC SCHOOL       64  0.854852
7   8               Bishop Shanahan        7  544449          BISHOP SHANAHAN HIGH SCHOOL       84  0.744183
8   9  Bishop Hendricken (American)        8  546469        BISHOP HENDRICKEN HIGH SCHOOL       87  0.736557
9  10                 Bishop Feehan        9  501234                        BISHOP FEEHAN       41  1.000000

linking_with_embeddings.R

library(tidyverse)
library(openai)


#-------------------------------------------------------------------------------
# Step 1: Open AI API key
#-------------------------------------------------------------------------------
# Your OpenAI API key should be an environment variable you set in ~.Renviron
# ... you should never put your API key directly in your code!


#-------------------------------------------------------------------------------
# Step 2: Define functions for later use:
#-------------------------------------------------------------------------------
# Function 1: calculate the cosine similarity between two vectors
cosine_similarity <- function(x, y) {
  sum(x * y) / (sqrt(sum(x^2)) * sqrt(sum(y^2)))
}

# Function 2: compute the maximum cosine similarity and its index between a vector and a list of vectors
max_cosine_similarity_index <- function(x, y_list) {
    # Use map_dbl() to compute the cosine similarity with each vector in y_list
    cos_sim <- map_dbl(y_list, cosine_similarity, x = x)
    # Return a list with the maximum value and its index of cosine similarity
    list(max_cos_sim = max(cos_sim), max_cos_sim_index = which.max(cos_sim))
}


#-------------------------------------------------------------------------------
# Step 3: Read in the data
#
# Assume you have two data frames called df1 and df2
# Each has a field with a school name
# For example: df1 has columns "id" and "HS"
#              df2 has columns "schid" and "name"
#-------------------------------------------------------------------------------
df1 <- read_csv("test1.csv")
print(df1)

df2 <- read_csv("test2.csv")
print(df2, n=100)


#-------------------------------------------------------------------------------
# Step 4: Get the embeddings for the key in each dataset
#
# This may take some time (took about 5 minutes for about 130 API calls)
# Cost is $0.0001 / 1K tokens
#-------------------------------------------------------------------------------
# Get the embeddings for the school names in df1 using the model text-embedding-ada-002
df1$embedding <- map(df1$HS, function(x) create_embedding(input = x, model = "text-embedding-ada-002")$data$embedding[[1]])

# Get the embeddings for the school names in df2 using the same model
df2$embedding <- map(df2$name, function(x) create_embedding(input = x, model = "text-embedding-ada-002")$data$embedding[[1]])


#-------------------------------------------------------------------------------
# Step 5: For each primary key, find the maximally similar foreign key in the 
#         other dataset
#-------------------------------------------------------------------------------
# Get the school name with the most similar embeddings
df1$similarity        <- map(df1$embedding, max_cosine_similarity_index, y_list = df2$embedding) %>% map_dbl("max_cos_sim")
df1$argmax_similarity <- map(df1$embedding, max_cosine_similarity_index, y_list = df2$embedding) %>% map_int("max_cos_sim_index")
df1$matched_name      <- df2$name[df1$argmax_similarity]


#-------------------------------------------------------------------------------
# Step 6: Inspect match performance
#-------------------------------------------------------------------------------
df1 %>% select(HS,matched_name,similarity) %>% print

linking_with_embeddings.Rout

> library(tidyverse)

> library(openai)

> #-------------------------------------------------------------------------------
> # Step 1: Open AI API key
> #------------------------------------ .... [TRUNCATED] 

> # Function 2: compute the maximum cosine similarity and its index between a vector and a list of vectors
> max_cosine_similarity_index <- function(x .... [TRUNCATED] 

> #-------------------------------------------------------------------------------
> # Step 3: Read in the data
> #
> # Assume you have two data frame .... [TRUNCATED] 
Rows: 15 Columns: 2                                                                                                                                               
── Column specification ───────────────────────────────────────────────────────────────────────────────────────────────────────
Delimiter: ","
chr (1): HS
dbl (1): id

ℹ Use `spec()` to retrieve the full column specification for this data.
ℹ Specify the column types or set `show_col_types = FALSE` to quiet this message.

> print(df1)
# A tibble: 15 × 2
      id HS                                       
   <dbl> <chr>                                    
 1     1 Timpanogos High School                   
 2     2 Timpanogos HS                            
 3     3 Timpanogos                               
 4     4 Archbishop Mitty                         
 5     5 Archbishop McCarthy (Fla.)               
 6     6 Archbishop Carroll                       
 7     7 Bishop Ryan Catholic                     
 8     8 Bishop Shanahan                          
 9     9 Bishop Hendricken (American)             
10    10 Bishop Feehan                            
11    11 Archbishop Molloy                        
12    12 Msgr. Bonner & Archbishop Prendergast    
13    13 Archbishop Spalding                      
14    14 Bishop O'Connell                         
15    15 Archbishop Edward A. McCarthy High School

> df2 <- read_csv("test2.csv")
Rows: 129 Columns: 2                                                                                                                                              
── Column specification ───────────────────────────────────────────────────────────────────────────────────────────────────────
Delimiter: ","
chr (1): name
dbl (1): schid

ℹ Use `spec()` to retrieve the full column specification for this data.
ℹ Specify the column types or set `show_col_types = FALSE` to quiet this message.

> print(df2, n=100)
# A tibble: 129 × 2
     schid name                                         
     <dbl> <chr>                                        
  1 100000 TIMPANOGOS HIGH                              
  2 111111 TIMPVIEW HIGH                                
  3 121212 TIMPSON H S                                  
  4 123456 ARCHBISHOP HANNA HIGH SCHOOL                 
  5 135792 ARCHBISHOP MITTY HIGH SCHOOL                 
  6 151515 ARCHBISHOP RIORDAN HIGH SCHOOL               
  7 161616 BISHOP ALEMANY HIGH SCHOOL                   
  8 171717 BISHOP AMAT MEMORIAL HIGH SCHOOL             
  9 181818 BISHOP CONATY-OUR LADY OF LORETTO HIGH SCHOOL
 10 191919 BISHOP GARCIA DIEGO HIGH SCHOOL              
 11 200000 BISHOP O DOWD HIGH SCHOOL                    
 12 222222 GRACE BISHOP SCHOOL                          
 13 234567 THE BISHOP'S SCHOOL                          
 14 246810 PATHWAY SCHOOL                               
 15 252525 ARCHBISHOP MCCARTHY HIGH SCHOOL              
 16 262626 BISHOP JOHN J SNYDER HIGH SCHOOL             
 17 272727 BISHOP KENNY HIGH SCHOOL                     
 18 282828 BISHOP MCLAUGHLIN CATHOLIC HIGH SCHOOL       
 19 292929 BISHOP VEROT CATHOLIC HIGH SCHOOL            
 20 300000 BISHOP KELLY HIGH SCHOOL                     
 21 323232 BISHOP CHATARD                               
 22 333333 BISHOP DWENGER HIGH SCHOOL                   
 23 345678 BISHOP LUERS HIGH SCHOOL                     
 24 353535 BISHOP GARRIGAN HIGH SCHOOL                  
 25 363636 BISHOP HEELAN CATHOLIC HIGH SCHOOL           
 26 373737 BISHOP HEELAN CATHOLIC SCHOOLS               
 27 383838 BISHOP CARROLL CATHOLIC HIGH SCHOOL          
 28 393939 BISHOP MIEGE HIGH SCHOOL                     
 29 400000 BISHOP SEABURY ACADEMY                       
 30 424242 BISHOP WARD HIGH SCHOOL                      
 31 444444 BISHOP BROSSART HIGH SCHOOL                  
 32 454545 ARCHBISHOP CHAPELLE HIGH SCHOOL              
 33 456789 ARCHBISHOP HANNAN HIGH SCHOOL                
 34 464646 ARCHBISHOP RUMMEL HIGH SCHOOL                
 35 474747 ARCHBISHOP SHAW HIGH SCHOOL                  
 36 484848 ARCHBISHOP CURLEY HIGH SCHOOL                
 37 494949 ARCHBISHOP SPALDING HIGH SCHOOL              
 38 500000 BISHOP MCNAMARA HIGH SCHOOL                  
 39 500009 BISHOP WALSH SCHOOL                          
 40 500129 ARCHBISHOP WILLIAMS HIGH SCHOOL              
 41 501019 BISHOP CONNOLLY HIGH SCHOOL                  
 42 501234 BISHOP FEEHAN                                
 43 501239 BISHOP FENWICK HIGH SCHOOL                   
 44 502029 SACRED HEART SCHOOL                          
 45 505059 BISHOP FOLEY CATHOLIC HIGH SCHOOL            
 46 506069 ARCHBISHOP BERGAN CATHOLIC SCHOOL            
 47 507079 BISHOP NEUMANN HIGH SCHOOL                   
 48 508089 BISHOP GORMAN HIGH SCHOOL                    
 49 509099 BISHOP MANOGUE CATHOLIC HIGH SCHOOL          
 50 510009 BISHOP BRADY HIGH SCHOOL                     
 51 510101 BISHOP GUERTIN HIGH SCHOOL                   
 52 511119 BISHOP EUSTACE PREP SCHOOL                   
 53 512129 BISHOP GEORGE AHR HIGH SCHOOL                
 54 512345 ARCHBISHOP MOLLOY HIGH SCHOOL                
 55 512349 ARCHBISHOP STEPINAC HIGH SCHOOL              
 56 513579 BISHOP GRIMES JR-SR HIGH SCHOOL              
 57 515159 BISHOP KEARNEY HIGH SCHOOL                   
 58 516169 BISHOP KEARNEY HIGH SCHOOL                   
 59 517179 BISHOP LOUGHLIN MEMORIAL HIGH SCHOOL         
 60 518189 BISHOP LUDDEN JUNIOR SENIOR HIGH SCHOOL      
 61 519199 BISHOP MAGINN HIGH SCHOOL                    
 62 520009 BISHOP MCGANN-MERCY DIOCESAN HIGH SCHOOL     
 63 520202 BISHOP TIMON-ST JUDE HIGH SCHOOL             
 64 522229 BISHOP MCGUINNESS CATHOLIC HIGH SCHOOL       
 65 523459 BISHOP RYAN CATHOLIC SCHOOL                  
 66 524689 ARCH BISHOP MOELLER HIGH SCHOOL              
 67 525252 ARCHBISHOP ALTER HIGH SCHOOL                 
 68 525259 ARCHBISHOP HOBAN HIGH SCHOOL                 
 69 526269 ARCHBISHOP MCNICHOLAS HIGH SCHOOL            
 70 527279 BISHOP FENWICK HIGH SCHOOL                   
 71 528289 BISHOP HARTLEY HIGH SCHOOL                   
 72 529299 BISHOP HOFFMAN CATHOLIC SCHOOL               
 73 530009 BISHOP READY HIGH SCHOOL                     
 74 532329 BISHOP ROSECRANS HIGH SCHOOL                 
 75 533339 BISHOP WATTERSON HIGH SCHOOL                 
 76 534569 BISHOP KELLEY HIGH SCHOOL, INC.              
 77 535359 BISHOP MCGUINNESS CATHOLIC HIGH SCHOOL       
 78 536369 ARCHBISHOP RYAN HIGH SCHOOL                  
 79 537379 ARCHBISHOP WOOD HIGH SCHOOL                  
 80 538389 BISHOP CANEVIN HIGH SCHOOL                   
 81 539399 BISHOP CARROLL CATHOLIC HIGH SCHOOL          
 82 540009 BISHOP GUILFOYLE CATHOLIC HIGH SCHOOL        
 83 542429 BISHOP MC CORT CATHOLIC HIGH SCHOOL          
 84 543210 BISHOP MCDEVITT HIGH SCHOOL                  
 85 544449 BISHOP SHANAHAN HIGH SCHOOL                  
 86 545459 THE EPISCOPAL ACADEMY                        
 87 545679 MONSIGNOR BONNER & ARCHBISHOP PRENDERGAST HS 
 88 546469 BISHOP HENDRICKEN HIGH SCHOOL                
 89 547479 ROBERT E LEE ACADEMY                         
 90 548489 BISHOP DUNNE CATHOLIC SCHOOL                 
 91 549499 BISHOP IRETON HIGH SCHOOL                    
 92 550009 BISHOP SULLIVAN CATHOLIC HIGH SCHOOL         
 93 550505 ARCHBISHOP MURPHY HIGH SCHOOL                
 94 552529 BISHOP BLANCHET HIGH SCHOOL                  
 95 554329 ARCHBISHOP COLEMAN CARROLL HIGH SCHOOL       
 96 555555 ARCHBISHOP CURLEY NOTRE DAME HIGH SCHOOL     
 97 555559 BISHOP MCNAMARA CATHOLIC HIGH SCHOOL         
 98 556569 BISHOP NOLL INSTITUTE                        
 99 556789 BISHOP LEBLOND HIGH SCHOOL                   
100 557579 ARCHBISHOP WALSH HIGH SCHOOL                 
# ℹ 29 more rows
# ℹ Use `print(n = ...)` to see more rows

> #-------------------------------------------------------------------------------
> # Step 4: Get the embeddings for the key in each dataset
> #
> #  .... [TRUNCATED] 

> # Get the embeddings for the school names in df2 using the same model
> df2$embedding <- map(df2$name, function(x) create_embedding(input = x, mo .... [TRUNCATED] 

> #-------------------------------------------------------------------------------
> # Step 5: For each primary key, find the maximally similar foreig .... [TRUNCATED] 

> df1$argmax_similarity <- map(df1$embedding, max_cosine_similarity_index, y_list = df2$embedding) %>% map_int("max_cos_sim_index")

> df1$matched_name      <- df2$name[df1$argmax_similarity]

> #-------------------------------------------------------------------------------
> # Step 6: Inspect match performance
> #-------------------------- .... [TRUNCATED] 
# A tibble: 15 × 3
   HS                                        matched_name                                 similarity
   <chr>                                     <chr>                                             <dbl>
 1 Timpanogos High School                    TIMPANOGOS HIGH                                   0.941
 2 Timpanogos HS                             TIMPANOGOS HIGH                                   0.936
 3 Timpanogos                                TIMPANOGOS HIGH                                   0.909
 4 Archbishop Mitty                          ARCHBISHOP MITTY HIGH SCHOOL                      0.898
 5 Archbishop McCarthy (Fla.)                ARCHBISHOP MCCARTHY HIGH SCHOOL                   0.881
 6 Archbishop Carroll                        ARCHBISHOP CARROLL HIGH SCHOOL                    0.906
 7 Bishop Ryan Catholic                      BISHOP RYAN CATHOLIC SCHOOL                       0.939
 8 Bishop Shanahan                           BISHOP SHANAHAN HIGH SCHOOL                       0.907
 9 Bishop Hendricken (American)              BISHOP HENDRICKEN HIGH SCHOOL                     0.877
10 Bishop Feehan                             BISHOP FEEHAN                                     0.925
11 Archbishop Molloy                         ARCHBISHOP MOLLOY HIGH SCHOOL                     0.914
12 Msgr. Bonner & Archbishop Prendergast     MONSIGNOR BONNER & ARCHBISHOP PRENDERGAST HS      0.913
13 Archbishop Spalding                       ARCHBISHOP SPALDING HIGH SCHOOL                   0.910
14 Bishop O'Connell                          BISHOP OCONNELL CATHOLIC HIGH SCHOOL              0.912
15 Archbishop Edward A. McCarthy High School ARCHBISHOP MCCARTHY HIGH SCHOOL                   0.927

test1.csv

test2.csv