Fuzzy Matching with Text Embedding Model

Overview

This repository explores using a fine-tuned T5 language model to resolve abbreviated, misspelled, or partially-written person names against a list of known candidate names, using embeddings derived from the model's encoder and cosine similarity. The core experiment (Fuzzy Matching.ipynb) fine-tunes t5-base on a small, hand-written synthetic dataset of "sentence completion" examples where shortened/garbled name mentions (e.g. "Z F and K P") must be expanded to full names (e.g. "Zikun Fu and Ken Pu"), then uses the fine-tuned model's encoder hidden states as name embeddings to rank candidate names by cosine similarity. Two related notebooks in the repo explore adjacent sub-problems: Entity Extraction.ipynb fine-tunes the same t5-base architecture to tag entity spans in sentences (a possible preprocessing step for locating which tokens to fuzzy-match), and Bird.ipynb is a separate, non-ML, rule-based utility that parses SQL queries and masks/classifies entities (tables, columns, aliases, literal values) using sqlparse and regex.

Repository contents

• Fuzzy Matching.ipynb — main fuzzy-matching-via-embeddings experiment (12 cells).
• Entity Extraction.ipynb — T5 fine-tuning for entity-span tagging (9 cells).
• Bird.ipynb — rule-based SQL entity extraction/masking script (6 cells), not related to embeddings.
• train.json / train_tables.json — text-to-SQL data (question/SQL pairs and database schemas) consumed by Bird.ipynb.
• test.ipynb — empty notebook, no content.

Method

Fuzzy Matching.ipynb (main experiment)

• Model: t5-base, loaded via transformers.T5Tokenizer and transformers.T5ForConditionalGeneration, fine-tuned with the Hugging Face Trainer API for 20 epochs (TrainingArguments(num_train_epochs=20, per_device_train_batch_size=8, ...)).
• Data: a synthetic, hand-written dataset embedded directly in the notebook — 61 training examples and 21 test examples, each an {"input": ..., "output": ...} pair. Inputs give a bracketed list of candidate full names plus a sentence containing abbreviated/garbled versions of those names (e.g. "[Zikun Fu, Nick Yang, Ken Pu] Complete the sentence: Z F and K P should meet N.Y."); outputs are the sentence with full names restored.
• Embedding/similarity approach: after fine-tuning, get_embedding() runs text through the fine-tuned T5 encoder and mean-pools last_hidden_state across tokens to produce a fixed-size vector. compare() uses sklearn.metrics.pairwise.cosine_similarity to score an input phrase's embedding against the embeddings of each candidate name and prints the candidates sorted by similarity (highest first).
• Separately, the fine-tuned model is also used generatively via model.generate() (complete_sentence()) to directly produce the disambiguated sentence, independent of the embedding/cosine-similarity path.

Entity Extraction.ipynb

• Fine-tunes t5-base (same tokenizer/model classes as above) for 5 epochs to convert a sentence into a version where named-entity tokens are replaced with <ENT> (entity start) / <CONT> (entity continuation) markers.
• Data: the Hugging Face dataset rjac/kaggle-entity-annotated-corpus-ner-dataset (loaded with datasets.load_dataset), sampled down to 2,000 rows (df.sample(n=2000, random_state=42)). The notebook computes a 70/15/15 train/validation/test split with sklearn.model_selection.train_test_split (into train_df/val_data/test_data), but the subsequent tokenization cell builds train_dataset from the original, unsplit 2,000-row train_data variable rather than from the 70% train_df split, so the model is trained on all 2,000 examples while validation/test use the 300/300-row splits drawn from that same pool.

Bird.ipynb

• A non-ML, rule-based script that uses sqlparse and regex to walk the tokens of a SQL query and classify each name token as an Alias, Table, Column, or Value (using train_tables.json schema metadata to distinguish table vs. column names), then reconstructs a "masked" version of the query with placeholder tokens (<extra_id_0..3>). It reads train.json (9,428 question/SQL records, each with db_id, question, evidence, SQL fields) and train_tables.json (per-database schema info) as input data. This notebook does not use any embedding model.

Results

The notebooks record actual executed outputs rather than a held-out benchmark evaluation; the concrete numbers below are copied verbatim from those outputs.

Fuzzy Matching.ipynb: after 20 epochs of fine-tuning, the final logged training metrics were train_loss: 4.403097397089004, and the final per-epoch evaluation showed eval_loss: 0.05864345282316208 (epoch 20.0). On one worked example (candidate names ["Emily Brown", "Elsa Betty", "David Grass", "David Green", "Finny Blade", "Fiona Black"], input mentioning "Em B and David G ... with Fi Blac"), the model's complete_sentence() generation correctly output "Emily Brown and David Green are chatting with Fiona Black," but the separate cosine-similarity ranking over the same candidates returned:

Elsa Betty: 0.5508
David Grass: 0.5338
Fiona Black: 0.5298
Finny Blade: 0.4101
David Green: 0.2918
Emily Brown: 0.2574

Note the embedding-based cosine-similarity ranking did not put the correct matches (David Green, Emily Brown) at the top for this example, even though generation-based decoding got the sentence right — i.e. the mean-pooled encoder embedding used for similarity ranking underperformed direct generation on this single example.

Entity Extraction.ipynb: after 5 epochs on the 2,000-example sampled dataset, the final logged metrics were train_loss: 0.7479060190320015 and a final eval_loss: 0.007703553885221481 (epoch 5.0). Two qualitative generation examples are shown in the notebook; for input "Ken Pu is going to meet Zikun Fu and Chen Yang on Sunday morning." the model output was "ENT> CONT> is going to meet ENT> CONT> and ENT> CONT> on ENT> morning." — the model detects entity spans but does not reproduce the leading < character of each tag correctly.

Bird.ipynb produces illustrative (non-numeric) output: for a sample query against the movie_platform database it correctly lists extracted entities such as ratings/movies (Table), movie_title/movie_id/rating_timestamp_utc (Column), T1/T2 (Alias), and numeric literals (Value), and prints the corresponding masked SQL string.

How to run

There is no requirements.txt in this repository. Based on the imports used across the notebooks, install:

pip install pandas torch transformers datasets scikit-learn sqlparse jupyter

Then, from the repository root:

jupyter notebook

• Fuzzy Matching.ipynb: self-contained — its training/test data is defined inline in the notebook. Run cells top to bottom.
• Entity Extraction.ipynb: downloads its dataset automatically from the Hugging Face Hub (rjac/kaggle-entity-annotated-corpus-ner-dataset) on first run; requires internet access. Run cells top to bottom.
• Bird.ipynb: requires train.json and train_tables.json, which are already present in the repository root. Run cells top to bottom.
• test.ipynb is empty and has no content to run.

Fuzzy Matching.ipynb and Entity Extraction.ipynb (the two notebooks that load a model) select cuda automatically if a GPU is available (torch.device('cuda' if torch.cuda.is_available() else 'cpu')), otherwise they fall back to CPU. Bird.ipynb does not use PyTorch or a GPU — it is a pure Python/regex/ sqlparse script.