File size: 2,405 Bytes
bde2f3a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | #!/usr/bin/env python3
"""FAST batch ingest SIGMOD P&D dataset directly into Redis.
Skips the HTTP API bottleneck — writes Redis keys directly.
"""
import csv
import hashlib
import json
import os
import sys
sys.path.insert(0, "/root/backend")
os.chdir("/root/backend")
import redis.asyncio as redis
from app.crypto_embeddings import get_embedder
CSV_PATH = "/root/tools/dark-collection/sigmod-pnd/Data/Telegram/Labeled/pred_pump_message.csv"
REDIS_HOST = os.getenv("REDIS_HOST", "rmi-redis")
REDIS_PORT = int(os.getenv("REDIS_PORT", "6379"))
REDIS_PASSWORD = os.getenv("REDIS_PASSWORD", "")
REDIS_DB = int(os.getenv("REDIS_DB", "0"))
BATCH_SIZE = 500
async def ingest():
with open(CSV_PATH) as f:
reader = csv.DictReader(f, delimiter="\t")
rows = list(reader)
print(f"Loading {len(rows)} P&D messages into Redis...")
r = redis.Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_DB)
embedder = await get_embedder()
count = 0
pipe = r.pipeline()
pipe_ops = 0
for _i, row in enumerate(rows):
msg = row.get("message", "")[:1000]
channel = row.get("channel_id", "unknown")
date = row.get("date", "")
content = f"[SIGMOD P&D] Channel {channel} - {date}: {msg}"
doc_id = hashlib.sha256(f"sigmod_{channel}_{date}".encode()).hexdigest()[:16]
# Generate embedding
vec = await embedder.embed_query(content)
key = f"rag:known_scams:{doc_id}"
doc = json.dumps(
{
"content": content,
"vector": vec,
"metadata": {
"source": "sigmod_pnd_2023",
"channel_id": channel,
"date": date,
"scam_type": "pump_and_dump",
},
}
)
pipe.set(key, doc)
pipe.sadd("rag:idx:known_scams", doc_id)
pipe_ops += 2
if pipe_ops >= BATCH_SIZE:
await pipe.execute()
count += BATCH_SIZE // 2
print(f" Progress: {count}/{len(rows)}")
pipe = r.pipeline()
pipe_ops = 0
if pipe_ops > 0:
await pipe.execute()
count += pipe_ops // 2
print(f"Done: {count} docs ingested into known_scams (Redis direct)")
await r.close()
if __name__ == "__main__":
import asyncio
asyncio.run(ingest())
|