Move all data to own folder. Add run shell script. Experiment (unsuccessfully so far) with example retreival
This commit is contained in:
William Jeynes
@@ -0,0 +1,360 @@
|
||||
import { parse } from "csv-parse";
|
||||
import fs from "fs";
|
||||
import { pipeline, cos_sim } from "@huggingface/transformers";
|
||||
import bm25Factory from "wink-bm25-text-search";
|
||||
import nlp from "wink-nlp-utils";
|
||||
import { logger } from "../utils/logger";
|
||||
import readline from "readline";
|
||||
|
||||
const CSV_PATHS = [
|
||||
"../data/dev-eng.csv",
|
||||
"../data/train-eng.csv",
|
||||
];
|
||||
|
||||
const CACHE_PATH = "../data/csv.cache.json";
|
||||
|
||||
const JSONL_PATH = "../data/results.jsonl"
|
||||
|
||||
type EmbeddingCache = {
|
||||
rawtexts: string[];
|
||||
cleantexts: string[];
|
||||
embeddings: number[][];
|
||||
};
|
||||
|
||||
export type RetrievalItem = {
|
||||
id: string | number;
|
||||
rawtext: string;
|
||||
cleantext?: string;
|
||||
};
|
||||
|
||||
export type RankedResult = RetrievalItem & {
|
||||
denseScore: number;
|
||||
sparseScore: number;
|
||||
fusedScore: number;
|
||||
};
|
||||
|
||||
let csvRawtexts: string[] = [];
|
||||
let csvCleantexts: string[] = [];
|
||||
let csvEmbeddings: number[][] = [];
|
||||
let csvBM25: any = null;
|
||||
let csvLoaded = false;
|
||||
|
||||
let jsonlRawtexts: string[] = [];
|
||||
let jsonlCleantexts: string[] = [];
|
||||
let jsonlEmbeddings: number[][] = [];
|
||||
let jsonlBM25: any = null;
|
||||
let jsonlLoaded = false;
|
||||
|
||||
|
||||
logger.info("Loading embedding model...");
|
||||
const featureExtractor = await pipeline(
|
||||
"feature-extraction",
|
||||
"Xenova/all-MiniLM-L6-v2"
|
||||
);
|
||||
logger.info("Embedding model loaded");
|
||||
|
||||
//Cached entrypoints
|
||||
export async function rankNormalizedClaims(
|
||||
query: string,
|
||||
topK = 5
|
||||
): Promise<RankedResult[]> {
|
||||
await ensureNormalizedClaimCSVLoaded();
|
||||
|
||||
logger.info("Ranking from CSV cache...");
|
||||
|
||||
const queryEmbedding = await embedText(query);
|
||||
|
||||
const denseScores = csvEmbeddings.map((docEmbedding) =>
|
||||
cos_sim(docEmbedding, queryEmbedding)
|
||||
);
|
||||
|
||||
const sparseScores = computeSparseScores(query, csvBM25, csvRawtexts);
|
||||
|
||||
const fusedScores = reciprocalRankFusion([denseScores, sparseScores]);
|
||||
|
||||
const ranked = csvRawtexts
|
||||
.map((text, i) => ({
|
||||
id: i,
|
||||
rawtext: text,
|
||||
cleantext: csvCleantexts[i],
|
||||
denseScore: denseScores[i],
|
||||
sparseScore: sparseScores[i],
|
||||
fusedScore: fusedScores[i],
|
||||
}))
|
||||
.sort((a, b) => b.fusedScore - a.fusedScore);
|
||||
|
||||
logger.info("Ranking complete (CSV mode)");
|
||||
|
||||
return ranked.slice(0, topK);
|
||||
}
|
||||
|
||||
export async function rankExampleTriggerEvents(
|
||||
query: string,
|
||||
topK = 5
|
||||
): Promise<RankedResult[]> {
|
||||
await ensureExampleClaimJsonlLoaded();
|
||||
|
||||
logger.info("Ranking from JSONL cache...");
|
||||
|
||||
const queryEmbedding = await embedText(query);
|
||||
|
||||
const denseScores = jsonlEmbeddings.map((docEmbedding) =>
|
||||
cos_sim(docEmbedding, queryEmbedding)
|
||||
);
|
||||
|
||||
const sparseScores = computeSparseScores(query, jsonlBM25, jsonlRawtexts);
|
||||
|
||||
const fusedScores = reciprocalRankFusion([denseScores, sparseScores]);
|
||||
|
||||
const ranked = jsonlRawtexts
|
||||
.map((text, i) => ({
|
||||
id: i,
|
||||
rawtext: text,
|
||||
cleantext: jsonlCleantexts[i],
|
||||
denseScore: denseScores[i],
|
||||
sparseScore: sparseScores[i],
|
||||
fusedScore: fusedScores[i],
|
||||
}))
|
||||
.sort((a, b) => b.fusedScore - a.fusedScore);
|
||||
|
||||
logger.info("Ranking complete (JSONL mode)");
|
||||
|
||||
return ranked.slice(0, topK);
|
||||
}
|
||||
|
||||
//Dynamic Entrypoint
|
||||
export async function rankDynamically(
|
||||
query: string,
|
||||
items: RetrievalItem[],
|
||||
topK = 5
|
||||
): Promise<RankedResult[]> {
|
||||
logger.info("Ranking dynamically (no cache)...");
|
||||
|
||||
if (!items.length) return [];
|
||||
|
||||
const texts = items.map((i) => i.rawtext);
|
||||
|
||||
const queryEmbedding = await embedText(query);
|
||||
|
||||
const docEmbeddings = await Promise.all(
|
||||
texts.map((text) => embedText(text))
|
||||
);
|
||||
|
||||
const denseScores = docEmbeddings.map((embedding) =>
|
||||
cos_sim(embedding, queryEmbedding)
|
||||
);
|
||||
|
||||
const localBM25 = buildBM25(texts);
|
||||
|
||||
const sparseScores = computeSparseScores(query, localBM25, texts);
|
||||
|
||||
const fusedScores = reciprocalRankFusion([denseScores, sparseScores]);
|
||||
|
||||
const ranked = items
|
||||
.map((item, i) => ({
|
||||
...item,
|
||||
denseScore: denseScores[i],
|
||||
sparseScore: sparseScores[i],
|
||||
fusedScore: fusedScores[i],
|
||||
}))
|
||||
.sort((a, b) => b.fusedScore - a.fusedScore);
|
||||
|
||||
logger.info("Ranking complete (dynamic mode)");
|
||||
|
||||
return ranked.slice(0, topK);
|
||||
}
|
||||
|
||||
//CSV stuff
|
||||
async function ensureNormalizedClaimCSVLoaded(): Promise<void> {
|
||||
if (csvLoaded) return;
|
||||
|
||||
logger.info("Initializing CSV ranking mode...");
|
||||
|
||||
if (fs.existsSync(CACHE_PATH)) {
|
||||
logger.info("Loading CSV cache from disk...");
|
||||
|
||||
const raw = fs.readFileSync(CACHE_PATH, "utf-8");
|
||||
const cache: EmbeddingCache = JSON.parse(raw);
|
||||
|
||||
csvRawtexts = cache.rawtexts;
|
||||
csvCleantexts = cache.cleantexts;
|
||||
csvEmbeddings = cache.embeddings;
|
||||
|
||||
logger.info("Loaded %s cached embeddings", csvEmbeddings.length);
|
||||
} else {
|
||||
logger.warn("CSV cache not found. Building embeddings...");
|
||||
|
||||
const seen = new Set<string>();
|
||||
|
||||
for (const path of CSV_PATHS) {
|
||||
await processNormalizationCSV(path, seen);
|
||||
}
|
||||
|
||||
const cache: EmbeddingCache = {
|
||||
rawtexts: csvRawtexts,
|
||||
cleantexts: csvCleantexts,
|
||||
embeddings: csvEmbeddings,
|
||||
};
|
||||
|
||||
fs.writeFileSync(CACHE_PATH, JSON.stringify(cache));
|
||||
logger.info("Cache written (%s embeddings)", csvEmbeddings.length);
|
||||
}
|
||||
|
||||
csvBM25 = buildBM25(csvRawtexts);
|
||||
|
||||
csvLoaded = true;
|
||||
logger.info("CSV mode ready");
|
||||
}
|
||||
|
||||
async function processNormalizationCSV(
|
||||
path: string,
|
||||
seen: Set<string>
|
||||
): Promise<void> {
|
||||
logger.info("Processing CSV: %s", path);
|
||||
|
||||
const stream = fs.createReadStream(path).pipe(parse());
|
||||
|
||||
for await (const row of stream) {
|
||||
const text = row[0];
|
||||
if (!text || seen.has(text)) continue;
|
||||
|
||||
seen.add(text);
|
||||
|
||||
const embedding = await embedText(text);
|
||||
|
||||
csvRawtexts.push(text);
|
||||
csvCleantexts.push(row[1]);
|
||||
csvEmbeddings.push(embedding);
|
||||
|
||||
if (csvRawtexts.length % 100 === 0) {
|
||||
logger.info("Embedded %s documents...", csvRawtexts.length);
|
||||
}
|
||||
}
|
||||
|
||||
logger.info("Finished CSV: %s", path);
|
||||
}
|
||||
|
||||
async function ensureExampleClaimJsonlLoaded(): Promise<void> {
|
||||
if (jsonlLoaded) return;
|
||||
|
||||
logger.info("Initializing JSONL ranking...");
|
||||
//TODO: make jsonl parsins
|
||||
const stream = fs.createReadStream(JSONL_PATH);
|
||||
|
||||
const rl = readline.createInterface({
|
||||
input: stream,
|
||||
crlfDelay: Infinity,
|
||||
});
|
||||
|
||||
for await (const line of rl) {
|
||||
if (!line.trim()) continue; // skip empty lines
|
||||
|
||||
const row = JSON.parse(line);
|
||||
|
||||
const text = row.text;
|
||||
|
||||
const embedding = await embedText(text);
|
||||
|
||||
jsonlRawtexts.push(text);
|
||||
|
||||
jsonlCleantexts.push(row.output[0].content);
|
||||
jsonlEmbeddings.push(embedding);
|
||||
}
|
||||
|
||||
|
||||
jsonlBM25 = buildBM25(jsonlRawtexts);
|
||||
|
||||
jsonlLoaded = true;
|
||||
logger.info("JSONL ranking done");
|
||||
}
|
||||
|
||||
|
||||
async function embedText(text: string): Promise<number[]> {
|
||||
const output = await featureExtractor(text, {
|
||||
pooling: "mean",
|
||||
normalize: true,
|
||||
});
|
||||
|
||||
return Array.from(output.data as Float32Array);
|
||||
}
|
||||
|
||||
function buildBM25(texts: string[]) {
|
||||
logger.info("Building BM25 index (%s docs)...", texts.length);
|
||||
|
||||
const bm25 = bm25Factory();
|
||||
|
||||
bm25.defineConfig({
|
||||
fldWeights: { text: 1 },
|
||||
bm25Params: { k1: 1.2, b: 0.75 },
|
||||
});
|
||||
|
||||
bm25.definePrepTasks([
|
||||
nlp.string.lowerCase,
|
||||
nlp.string.tokenize0,
|
||||
nlp.tokens.removeWords,
|
||||
]);
|
||||
|
||||
texts.forEach((text, i) => {
|
||||
bm25.addDoc({ text }, i);
|
||||
});
|
||||
|
||||
bm25.consolidate();
|
||||
|
||||
logger.info("BM25 ready");
|
||||
|
||||
return bm25;
|
||||
}
|
||||
|
||||
function computeSparseScores(
|
||||
query: string,
|
||||
bm25: any,
|
||||
texts: string[]
|
||||
): number[] {
|
||||
const results = bm25.search(query);
|
||||
|
||||
const scores = new Array(texts.length).fill(0);
|
||||
|
||||
results.forEach((r: any) => {
|
||||
scores[r[0]] = r[1];
|
||||
});
|
||||
|
||||
return scores;
|
||||
}
|
||||
|
||||
function reciprocalRankFusion(
|
||||
scoreLists: number[][],
|
||||
k = 60
|
||||
): number[] {
|
||||
const length = scoreLists[0].length;
|
||||
const fused = new Array(length).fill(0);
|
||||
|
||||
for (const scores of scoreLists) {
|
||||
const ranked = scores
|
||||
.map((score, i) => ({ score, i }))
|
||||
.sort((a, b) => b.score - a.score)
|
||||
.map((x) => x.i);
|
||||
|
||||
ranked.forEach((docIndex, rank) => {
|
||||
fused[docIndex] += 1 / (k + rank);
|
||||
});
|
||||
}
|
||||
|
||||
return fused;
|
||||
}
|
||||
|
||||
// console.log(await rankFromCSV("barrack obama"))
|
||||
// console.log(
|
||||
// await rankDynamically(
|
||||
// "i fell over",
|
||||
// [
|
||||
// { id: 1, rawtext: "I slipped and fell on the floor." },
|
||||
// { id: 2, rawtext: "Barack Obama was the 44th president." },
|
||||
// { id: 3, rawtext: "He tripped and hurt his knee badly." },
|
||||
// { id: 4, rawtext: "The weather is sunny today." },
|
||||
// { id: 5, rawtext: "She lost her balance and fell down the stairs." },
|
||||
// ]
|
||||
// )
|
||||
// );
|
||||
|
||||
// await ensureExampleClaimJsonlLoaded()
|
||||
// console.log(await rankExampleTriggerEvents("Niger"))
|
||||
Reference in New Issue
Block a user