razsarusi's picture
Add co-author Tomer Dariel
d922e44 verified
|
Raw
History Blame Contribute Delete
4.47 kB
metadata
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

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 · App: razsarusi/plottwist