feat(ml): add TA-Lib annotation generation and import workflow
Add complete workflow for using TA-Lib to bootstrap training data: - generate_talib_annotations.py: Python script to run TA-Lib CDL* functions and output span annotations in UI-compatible format - import_talib_annotations.ts: TypeScript script to import generated annotations into the UI database with auto-label-type creation - npm script 'import-annotations' for easy execution - TALIB_WORKFLOW.md: Comprehensive guide covering the full cycle: * Generate patterns with TA-Lib * Import into UI * Review and edit in browser * Export and train model * Compare predictions with TA-Lib detections * Iterate for improvement This enables the intended workflow: use TA-Lib for initial annotations, manually refine them, then train a model that learns from corrections.
This commit is contained in:
parent
228f70daf3
commit
847ff67986
18 changed files with 5416 additions and 7 deletions
252
scripts/import_talib_annotations.ts
Normal file
252
scripts/import_talib_annotations.ts
Normal file
|
|
@ -0,0 +1,252 @@
|
|||
#!/usr/bin/env tsx
|
||||
/**
|
||||
* Import TA-Lib generated annotations into the Candle Annotator database.
|
||||
*
|
||||
* This script reads a JSON file with annotations (from generate_talib_annotations.py)
|
||||
* and imports them as span annotations that can be viewed and edited in the UI.
|
||||
*
|
||||
* Usage:
|
||||
* npm run import-annotations -- --file talib_annotations.json --chart-id 1
|
||||
*/
|
||||
|
||||
import { readFile } from 'fs/promises';
|
||||
import { db } from '../src/lib/db';
|
||||
import { spanAnnotations, spanLabelTypes } from '../src/lib/db/schema';
|
||||
import { eq } from 'drizzle-orm';
|
||||
|
||||
interface Annotation {
|
||||
start_time: number;
|
||||
end_time: number;
|
||||
label: string;
|
||||
confidence?: number;
|
||||
source?: string;
|
||||
notes?: string;
|
||||
}
|
||||
|
||||
interface ImportData {
|
||||
annotations: Annotation[];
|
||||
metadata?: {
|
||||
source?: string;
|
||||
count?: number;
|
||||
};
|
||||
}
|
||||
|
||||
interface LabelTypeMap {
|
||||
[key: string]: {
|
||||
id: number;
|
||||
color: string;
|
||||
};
|
||||
}
|
||||
|
||||
async function ensureLabelTypes(labels: string[]): Promise<LabelTypeMap> {
|
||||
/**
|
||||
* Ensure span label types exist for all unique labels.
|
||||
* Creates missing label types with auto-generated colors.
|
||||
*/
|
||||
const uniqueLabels = [...new Set(labels)];
|
||||
const labelMap: LabelTypeMap = {};
|
||||
|
||||
console.log(`\nEnsuring ${uniqueLabels.length} label types exist...`);
|
||||
|
||||
// Fetch existing label types
|
||||
const existing = await db.select().from(spanLabelTypes);
|
||||
const existingMap = new Map(existing.map(lt => [lt.name, lt]));
|
||||
|
||||
// Color palette for auto-generated labels
|
||||
const colors = [
|
||||
'#22c55e', // green (bullish)
|
||||
'#ef4444', // red (bearish)
|
||||
'#3b82f6', // blue
|
||||
'#f59e0b', // amber
|
||||
'#8b5cf6', // purple
|
||||
'#ec4899', // pink
|
||||
'#06b6d4', // cyan
|
||||
'#f97316', // orange
|
||||
];
|
||||
|
||||
let colorIndex = 0;
|
||||
let sortOrder = existing.length;
|
||||
|
||||
for (const label of uniqueLabels) {
|
||||
if (existingMap.has(label)) {
|
||||
const existing = existingMap.get(label)!;
|
||||
labelMap[label] = {
|
||||
id: existing.id,
|
||||
color: existing.color,
|
||||
};
|
||||
console.log(` ✓ ${label} (existing, id: ${existing.id})`);
|
||||
} else {
|
||||
// Assign color based on bullish/bearish
|
||||
let color: string;
|
||||
if (label.toLowerCase().includes('bullish')) {
|
||||
color = '#22c55e'; // green
|
||||
} else if (label.toLowerCase().includes('bearish')) {
|
||||
color = '#ef4444'; // red
|
||||
} else {
|
||||
color = colors[colorIndex % colors.length];
|
||||
colorIndex++;
|
||||
}
|
||||
|
||||
// Create new label type
|
||||
const [newLabel] = await db.insert(spanLabelTypes).values({
|
||||
name: label,
|
||||
display_name: label,
|
||||
color,
|
||||
hotkey: null,
|
||||
is_active: 1,
|
||||
sort_order: sortOrder++,
|
||||
}).returning();
|
||||
|
||||
labelMap[label] = {
|
||||
id: newLabel.id,
|
||||
color: newLabel.color,
|
||||
};
|
||||
|
||||
console.log(` + ${label} (created, id: ${newLabel.id}, color: ${color})`);
|
||||
}
|
||||
}
|
||||
|
||||
return labelMap;
|
||||
}
|
||||
|
||||
async function importAnnotations(
|
||||
annotations: Annotation[],
|
||||
chartId: number,
|
||||
labelMap: LabelTypeMap
|
||||
): Promise<number> {
|
||||
/**
|
||||
* Import annotations into the database.
|
||||
* Returns the number of annotations imported.
|
||||
*/
|
||||
console.log(`\nImporting ${annotations.length} annotations for chart ${chartId}...`);
|
||||
|
||||
let imported = 0;
|
||||
let skipped = 0;
|
||||
|
||||
for (const ann of annotations) {
|
||||
const labelInfo = labelMap[ann.label];
|
||||
if (!labelInfo) {
|
||||
console.error(` ✗ Skipping: Unknown label "${ann.label}"`);
|
||||
skipped++;
|
||||
continue;
|
||||
}
|
||||
|
||||
try {
|
||||
await db.insert(spanAnnotations).values({
|
||||
chart_id: chartId,
|
||||
start_time: ann.start_time,
|
||||
end_time: ann.end_time,
|
||||
label: ann.label,
|
||||
confidence: ann.confidence || null,
|
||||
outcome: null,
|
||||
notes: ann.notes || null,
|
||||
sub_spans: null,
|
||||
color: labelInfo.color,
|
||||
source: ann.source || 'programmatic',
|
||||
model_prediction: null,
|
||||
});
|
||||
|
||||
imported++;
|
||||
|
||||
if (imported % 100 === 0) {
|
||||
console.log(` ${imported}/${annotations.length} imported...`);
|
||||
}
|
||||
} catch (error: any) {
|
||||
console.error(` ✗ Error importing annotation: ${error.message}`);
|
||||
skipped++;
|
||||
}
|
||||
}
|
||||
|
||||
console.log(`\n✓ Imported ${imported} annotations`);
|
||||
if (skipped > 0) {
|
||||
console.log(`✗ Skipped ${skipped} annotations`);
|
||||
}
|
||||
|
||||
return imported;
|
||||
}
|
||||
|
||||
async function main() {
|
||||
const args = process.argv.slice(2);
|
||||
|
||||
// Parse arguments
|
||||
let filePath: string | null = null;
|
||||
let chartId: number | null = null;
|
||||
let clearExisting = false;
|
||||
|
||||
for (let i = 0; i < args.length; i++) {
|
||||
if (args[i] === '--file' || args[i] === '-f') {
|
||||
filePath = args[++i];
|
||||
} else if (args[i] === '--chart-id' || args[i] === '-c') {
|
||||
chartId = parseInt(args[++i], 10);
|
||||
} else if (args[i] === '--clear') {
|
||||
clearExisting = true;
|
||||
} else if (args[i] === '--help' || args[i] === '-h') {
|
||||
console.log(`
|
||||
Import TA-Lib annotations into Candle Annotator database
|
||||
|
||||
Usage:
|
||||
npm run import-annotations -- --file <json-file> --chart-id <id> [--clear]
|
||||
|
||||
Options:
|
||||
--file, -f <path> Input JSON file (from generate_talib_annotations.py)
|
||||
--chart-id, -c <id> Chart ID to import annotations into
|
||||
--clear Clear existing annotations for this chart before import
|
||||
--help, -h Show this help message
|
||||
|
||||
Example:
|
||||
npm run import-annotations -- --file talib_annotations.json --chart-id 1
|
||||
`);
|
||||
process.exit(0);
|
||||
}
|
||||
}
|
||||
|
||||
if (!filePath || chartId === null) {
|
||||
console.error('Error: --file and --chart-id are required');
|
||||
console.error('Run with --help for usage information');
|
||||
process.exit(1);
|
||||
}
|
||||
|
||||
console.log('=== TA-Lib Annotation Import ===\n');
|
||||
console.log(`Input file: ${filePath}`);
|
||||
console.log(`Chart ID: ${chartId}`);
|
||||
console.log(`Clear existing: ${clearExisting ? 'yes' : 'no'}`);
|
||||
|
||||
// Read annotations file
|
||||
console.log('\nReading annotations file...');
|
||||
const fileContent = await readFile(filePath, 'utf-8');
|
||||
const data: ImportData = JSON.parse(fileContent);
|
||||
|
||||
console.log(`Found ${data.annotations.length} annotations`);
|
||||
if (data.metadata) {
|
||||
console.log(`Source: ${data.metadata.source || 'unknown'}`);
|
||||
}
|
||||
|
||||
// Clear existing annotations if requested
|
||||
if (clearExisting) {
|
||||
console.log('\nClearing existing annotations...');
|
||||
const result = await db.delete(spanAnnotations)
|
||||
.where(eq(spanAnnotations.chart_id, chartId));
|
||||
console.log(`Deleted existing annotations`);
|
||||
}
|
||||
|
||||
// Collect unique labels
|
||||
const uniqueLabels = [...new Set(data.annotations.map(a => a.label))];
|
||||
|
||||
// Ensure label types exist
|
||||
const labelMap = await ensureLabelTypes(uniqueLabels);
|
||||
|
||||
// Import annotations
|
||||
const imported = await importAnnotations(data.annotations, chartId, labelMap);
|
||||
|
||||
console.log('\n=== Import Complete ===');
|
||||
console.log(`\nNext steps:`);
|
||||
console.log(`1. Open http://localhost:3000 and select chart ${chartId}`);
|
||||
console.log(`2. Review the TA-Lib generated annotations`);
|
||||
console.log(`3. Edit, delete, or add new annotations as needed`);
|
||||
console.log(`4. Export and train: npm run ml:export-and-train`);
|
||||
}
|
||||
|
||||
main().catch((error) => {
|
||||
console.error('\n✗ Import failed:', error.message);
|
||||
process.exit(1);
|
||||
});
|
||||
Loading…
Add table
Add a link
Reference in a new issue