baya1116's picture
Super-squash branch 'main' using huggingface_hub
b5989f0
Raw
History Blame Contribute Delete
10.6 kB
"""Train a small intent classifier on BGE-small embeddings to replace the brittle regex gates
(`_is_question` / `_is_factlike` / `classify`). 3 classes:
- question : the user wants a specific value recalled or looked up -> RETRIEVE
- fact : the user asserts a value/attribute (incl. corrections) -> LOG to memory
- chitchat : greeting / acknowledgement / open-ended generation -> just respond
Frozen BGE-small embeddings + sklearn LogisticRegression. Saves evals/intent_clf.joblib.
Run: python3.12 evals/intent_train.py
"""
import sys, os, json
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "runtime"))
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
import numpy as np
DATA = {
"recall": [ # a question about the USER's own info -> retrieve personal memory (L1/L2)
"what's my insurance policy number?", "where did I park?", "which spot did I leave the car in",
"how many guests are coming now", "how many people are coming", "remind me the theme of the party",
"who is the party for", "what flavor is the cake", "tell me my employee id", "do you remember my address",
"what was my reservation code again", "can you recall my locker combination", "what's the code for my gym locker",
"what time is my flight", "when is the meeting",
"what did I say my budget was", "how much did the goodie bags cost", "what's the total again",
"remind me what color I chose", "which hotel am I staying at", "what's my wifi password",
"what's the name of my project", "how old is my daughter", "what's my blood type",
"what's my seat number", "recall my emergency contact", "what's my account balance",
"what color am I painting it and how many shelves does it have now", "where did I leave the vehicle",
"is there any food I should avoid", "what dosage did you mention", "what city did I say I'm from",
"did I mention my dog's name", "whats my flight reservation code", "is my appointment confirmed",
"what's the new deadline", "remind me my gym locker code", "what was the venue again", "how many shelves now",
"what's my employee ID", "what is my name", "tell me where my keys are", "what's the policy number",
"how long am I staying", "what dietary preference should I mention", "what's my reservation under",
],
"lookup": [ # a question about a GENERAL / WORLD fact -> web search
"who is the current ceo of openai", "how tall is mount fuji", "what is the capital of bhutan",
"which country won the 2022 world cup", "who wrote neuromancer", "what's the population of reykjavik",
"who is the current emperor of japan", "what's the capital of france", "when did world war 2 end",
"who painted the mona lisa", "what's the tallest building in the world", "who is the prime minister of japan",
"how far is the moon from earth", "what's the speed of light", "who discovered penicillin",
"what year did the berlin wall fall", "what's the largest ocean", "who invented the telephone",
"what's the boiling point of water", "how many continents are there", "what currency is used in japan",
"who is the ceo of tesla", "what's the national bird of the usa", "when was the eiffel tower built",
],
"fact": [
"I'll paint it blue", "the theme is space", "there will be 8 guests",
"actually, 2 more are coming, so 10 now", "make it 6 shelves instead of 5", "change the theme to dinosaurs",
"my car is parked in bay 12", "I'm allergic to peanuts", "my gym locker combination is 5588",
"I'll order a chocolate cake", "my name is Aki", "my insurance policy number is POL-55821",
"I live in Sapporo", "my flight is at 3pm", "the meeting moved to Friday", "I prefer window seats",
"my budget is $500", "my dog is named Mochi", "I work in the Helsinki office", "my employee id is EMP-90832",
"remember that my wifi password is hunter2", "note that I'm vegetarian", "my daughter is 7 years old",
"I'm staying at the Hilton", "the project is called Apollo", "my reservation code is QX7-2291",
"I just parked on level B3, spot 47", "my address is 42 Oak Street", "the password changed to abc123",
"my blood type is O negative", "I drive a red Toyota", "the deadline is next Monday", "my seat is 14C",
"I'm bringing 3 bottles of wine", "the venue is downtown", "scratch that, make it 4 instead",
"my emergency contact is my sister Mei", "I'd like the room painted white", "the cake should be gluten free",
"set the guest count to 12", "my new phone number is 555-0199", "we're meeting at the cafe instead",
"the color should be matte black", "my locker is number 77", "it costs five dollars each",
"the party is for my friend Mia", "each shelf is 80 cm wide", "I'll have the salmon",
]
,
"math": [
"what's 15% of 240", "if I buy 3 apples at $2 each, what's the total", "how much is 17 times 24",
"I have twelve cookies and eat five, how many are left", "what's 1000 plus 250",
"split 80 dollars equally among 4 people", "convert 3.5 hours to minutes",
"if a train goes 60 km in 1.5 hours, how far in 4 hours", "a $50 item is 20% off, what's the final price",
"what's the sum of 8, 13, and 21", "how many minutes in 2 and a half hours",
"if each bag costs 5 dollars and there are 8 guests, total cost?", "10 percent of 350 is what",
"divide 144 by 12", "three times seven plus two", "what is 1.05 cubed times 1000",
"I worked 8 hours at 15 an hour, how much did I earn", "round 3776.24 to the nearest hundred",
"nine times five", "what's twelve minus five", "add seven and thirteen", "subtract five from twelve",
"what is fifteen percent of two hundred", "twelve minus five equals what", "half of forty",
"three hundred divided by twelve", "double sixteen", "sum of eight, thirteen and twenty",
"what's seven times eight", "ninety minus forty-two", "a dozen plus five", "two thirds of ninety",
"how much is twenty percent of fifty", "if I have twelve apples and eat five, how many are left",
"what's the total of nine and six", "subtract nineteen from a hundred", "five squared",
],
"command": [
"recompute the total", "redo that calculation", "recalculate with the new number",
"what's the new total now", "do the math again", "update the estimate", "add one more and recompute",
"recalc the budget", "figure it out again with 10 instead", "adjust the total for the change",
"redo it", "compute it again please", "give me the updated figure", "recalculate the cost",
"now total it up again", "work out the new amount", "redo the sum with the correction",
],
"chitchat": [
"hey, good morning!", "thanks so much", "explain compound interest briefly", "tell me a joke",
"how are you doing today", "that's great, thanks", "can you help me plan a party", "let's get started",
"write me a short poem", "what can you help me with", "explain how a hash map works",
"give me some ideas for dinner", "sounds good", "haha nice", "ok cool", "summarize this for me",
"draft an email to my landlord", "good night", "I appreciate it", "let's chat about something else",
"describe the water cycle", "brainstorm some hobby ideas", "yes please", "no thanks",
"tell me about black holes", "make it a little more firm", "that works for me", "perfect, thank you",
"could you explain it differently", "i'm not sure what to do", "let's switch topics",
"I'm planning a weekend trip", "I'm thinking about redecorating", "I'm looking for some ideas",
"I'm trying to decide what to cook", "I want to get into a new hobby", "let's plan something fun",
"explain what a ryokan is", "what's a ryokan", "tell me what a mutex is", "what does compounding mean",
"explain what an API is", "describe what a black hole is", "what is machine learning in simple terms",
"explain how compound interest works", "describe what inflation means", "walk me through how mortgages work",
"teach me about budgeting basics", "give me an overview of interest rates", "explain probability simply",
"describe how percentages work", "explain the idea behind compound growth", "tell me about the stock market briefly",
"give a simple explanation of averages", "explain statistics in plain terms", "walk me through how taxes work",
"walk me through the steps", "anything else you'd suggest", "hello there", "got it",
"please continue", "never mind", "that's hilarious", "do you think that's a good idea",
"help me write a tweet", "give me a recipe for pasta", "what's a good beginner hobby",
],
}
def main():
from rag import BGERetriever
bge = BGERetriever()
items_by_label = {k: list(v) for k, v in DATA.items()} # hand-labelled core
gen_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "intent_gen.jsonl")
ngen = 0
if os.path.exists(gen_path) and "--hand-only" not in sys.argv: # + open-model-generated examples
import json
have = {t for v in items_by_label.values() for t in v}
for l in open(gen_path):
o = json.loads(l)
if o["label"] in items_by_label and o["text"] not in have:
items_by_label[o["label"]].append(o["text"]); have.add(o["text"]); ngen += 1
print(f"merged {ngen} generated examples")
X, y, texts = [], [], []
for label, items in items_by_label.items():
X.append(bge._encode(items, is_query=False)) # PURE BGE-small (384-d) — no hand
y += [label] * len(items); texts += items # syntactic features (data suffices now)
X = np.concatenate(X, 0)
y = np.array(y)
print(f"dataset: {len(y)} examples | {[ (l, int((y==l).sum())) for l in DATA ]}")
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import StratifiedKFold, cross_val_predict
from sklearn.metrics import classification_report, confusion_matrix
clf = LogisticRegression(max_iter=2000, C=2.0, class_weight="balanced")
# honest held-out estimate via 5-fold CV
pred = cross_val_predict(clf, X, y, cv=StratifiedKFold(5, shuffle=True, random_state=0))
print("\n5-fold CV report:\n", classification_report(y, pred, digits=3))
print("confusion (rows=true q/f/c):\n", confusion_matrix(y, pred, labels=list(DATA)))
# train final on all data, save
clf.fit(X, y)
import joblib
out = os.path.join(os.path.dirname(os.path.abspath(__file__)), "intent_clf.joblib")
joblib.dump({"clf": clf, "labels": list(DATA)}, out)
print("saved", out)
# show a few CV mistakes (where the regex/embedding boundary is fuzzy)
print("\nmisclassified:")
for t, yt, yp in zip(texts, y, pred):
if yt != yp:
print(f" {yt}->{yp}: {t}")
print("INTENT_TRAIN_DONE")
if __name__ == "__main__":
main()