razsarusi's picture
Add co-author Tomer Dariel
d922e44 verified
|
Raw
History Blame Contribute Delete
4.47 kB
---
license: apache-2.0
library_name: sentence-transformers
pipeline_tag: sentence-similarity
base_model: sentence-transformers/all-mpnet-base-v2
tags:
- sentence-transformers
- feature-extraction
- movie-retrieval
- fine-tuned
---
# PlotTwist Embedder (fine-tuned)
**Authors:** Raz Sarusi Β· Tomer Dariel
The retrieval model behind **PlotTwist**: it maps a short movie *pitch* to the closest *loglines* in a
10,000-row synthetic catalog. This repo hosts the **fine-tuned** embedder plus the FAISS index and the
catalog, so the app loads everything at startup with no re-encoding.
## 1. Model selection β€” we compared **three** HF embedding models
Before any fine-tuning we embedded the full 10k catalog with three models and compared them on
**retrieval quality**, **encode time**, and **vector size** (full live table in `Part3_Recommendation.ipynb`):
| model | dim | notes | role |
|---|---|---|---|
| `sentence-transformers/all-MiniLM-L6-v2` | 384 | small & fast | speed baseline |
| `sentence-transformers/all-mpnet-base-v2` | 768 | strongest retrieval quality | **winner** |
| `BAAI/bge-small-en-v1.5` | 384 | strong small model | small-model contender |
Selection rule: **primary = retrieval quality**, ties broken by faster encoding / smaller size.
`all-mpnet-base-v2` won on quality, so it became the base we fine-tuned.
## 2. Fine-tuning + a non-circular evaluation
v1's only metric β€” `genre-consistency@k` β€” is **circular** (data was generated conditioned on genre,
then genre agreement was measured). v2 replaces it with a genuine held-out task:
1. A small LLM writes the short **user-style pitch** for a sample of loglines β†’ `(pitch, logline)` pairs.
2. **recall@1 / @5 / @10**: does a test pitch retrieve its *true* logline from the full 10k catalog?
3. Fine-tune with `MultipleNegativesRankingLoss` (in-batch negatives) on the `(pitch, logline)` pairs.
**Results** (held-out test = 600 pitches, catalog = 10k, train/test = 2400/600):
| metric | base | fine-tuned | Ξ” |
|---|---|---|---|
| recall@1 | 0.753 | **0.845** | **+0.092** |
| recall@5 | 0.890 | **0.933** | **+0.043** |
| recall@10 | 0.918 | **0.953** | **+0.035** |
recall@1 improved **+9.2 points (+12.2% relative)** β€” a real, non-circular improvement.
### Training configuration (the training "matrix")
| setting | value |
|---|---|
| base model | `sentence-transformers/all-mpnet-base-v2` (768-d, cosine) |
| objective / loss | `MultipleNegativesRankingLoss` (scale 20.0, `cos_sim`) |
| training pairs | 2,400 `(pitch β†’ logline)`; held-out test 600 |
| epochs Β· batch size | 2 Β· 32 (β‰ˆ 150 optimization steps) |
| learning rate Β· schedule | 5e-5 Β· linear Β· AdamW (fused) |
| seed Β· training time | 42 Β· β‰ˆ 1.9 min Β· Sentence-Transformers 5.6.0 |
**Training loss vs. evaluation metric β€” and why the recall@k is the real proof.**
`MultipleNegativesRankingLoss` is a *contrastive ranking* loss: its absolute value is not an accuracy
score, and the run was short (β‰ˆ150 steps), so a per-step loss curve is not the meaningful signal here.
The rigorous evidence that fine-tuning **beat the base model** is the held-out **recall@k lift** in the
table above β€” measured on pitches the model never trained on. Over the 2 epochs the model learned to
place each pitch next to its true logline, moving **recall@1 from 0.753 β†’ 0.845** (base β†’ fine-tuned).
## 3. Files
- fine-tuned `all-mpnet-base-v2` weights (Sentence-Transformers format)
- `plottwist.faiss` β€” FAISS `IndexFlatIP` over L2-normalized embeddings (= cosine), 10k vectors
- `plottwist_catalog.parquet` β€” the catalog rows aligned to the index
- `eval_recall_base_vs_ft.json` β€” the numbers above
## 4. Usage
```python
from sentence_transformers import SentenceTransformer
import faiss, pandas as pd
from huggingface_hub import hf_hub_download
model = SentenceTransformer("razsarusi/plottwist-embedder-ft")
idx = faiss.read_index(hf_hub_download("razsarusi/plottwist-embedder-ft", "plottwist.faiss"))
cat = pd.read_parquet(hf_hub_download("razsarusi/plottwist-embedder-ft", "plottwist_catalog.parquet"))
q = model.encode(["a lonely lighthouse keeper bargains with the sea"], normalize_embeddings=True)
D, I = idx.search(q.astype("float32"), 3)
print(cat.iloc[I[0]][["title","genre","logline"]])
```
Dataset: [`razsarusi/plottwist-movies`](https://huggingface.co/datasets/razsarusi/plottwist-movies) Β·
App: [`razsarusi/plottwist`](https://huggingface.co/spaces/razsarusi/plottwist)