Spaces:
Running
Running
File size: 4,347 Bytes
122cc3c | 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 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | """Generate extra mentor rows to supplement the clean dataset."""
import asyncio
import json
import os
import random
import sys
from pathlib import Path
import aiohttp
from dotenv import load_dotenv
load_dotenv()
API_KEY = os.getenv("API_KEY") or os.getenv("ZENMUX_API_KEY")
BASE_URL = os.getenv("BASE_URL") or os.getenv("ZENMUX_BASE_URL", "https://api.tokenrouter.com/v1")
MODEL = os.getenv("MODEL") or os.getenv("ZENMUX_MODEL", "MiniMax-M3")
ROOT = Path(__file__).resolve().parent.parent
CLEAN_FILE = ROOT / "data" / "retro-alpha-clean.jsonl"
EXTRA_FILE = ROOT / "data" / "retro-alpha-mentor-extra.jsonl"
ASSETS = ["cash", "fd", "gov_bonds", "nifty_50", "nifty_it", "real_estate", "crypto", "gold"]
STYLES = ["savage", "educational", "sarcastic", "encouraging"]
headers = {"Authorization": f"Bearer {API_KEY}", "Content-Type": "application/json"}
def random_portfolio() -> dict:
weights = [random.random() for _ in ASSETS]
total = sum(weights)
return {a: round(w / total, 3) for a, w in zip(ASSETS, weights)}
def build_prompt(style: str) -> dict:
allocations = random_portfolio()
if random.random() < 0.6:
heavy = random.choice(["crypto", "nifty_it", "real_estate"])
allocations = {a: 0.05 for a in ASSETS}
allocations[heavy] = 0.65
start = random.choice([500000, 1000000, 2000000, 5000000])
drawdown = round(random.uniform(-0.05, -0.45), 2)
end = int(start * (1 + drawdown))
sharpe = round(random.uniform(-1.2, 1.5), 2)
system = "You are a finance professor NPC in a video game. Output a year-end review in this exact short format:\nroast: <roast, max 50 chars>\nsharpe_ratio: <number>\nlesson: <lesson, max 80 chars>\nsuggestion: <tip, max 50 chars>"
user = f"""Example:
roast: 65% crypto? Portfolio vanished.
sharpe_ratio: -0.5
lesson: Sharpe measures return per unit risk.
suggestion: Cap crypto at 10%.
Now generate for: start=₹{start/100000:.1f}L, end=₹{end/100000:.1f}L, drawdown={drawdown*100:.0f}%, sharpe={sharpe}, style={style}"""
return {"system": system, "user": user, "task": "sharpe_mentor", "style": style}
def clean(text: str) -> str:
text = text.strip()
while "<think>" in text and "</think>" in text:
s = text.find("<think>")
e = text.find("</think>") + len("</think>")
text = text[:s] + text[e:]
return text.strip()
async def call_one(session, prompt, semaphore):
payload = {
"model": MODEL,
"messages": [
{"role": "system", "content": prompt["system"]},
{"role": "user", "content": prompt["user"]}
],
"temperature": 0.7,
"max_tokens": 500
}
for attempt in range(5):
async with semaphore:
try:
async with session.post(f"{BASE_URL}/chat/completions", headers=headers, json=payload, timeout=60) as resp:
data = await resp.json()
content = data["choices"][0]["message"].get("content", "")
if content.strip():
return {
"task": prompt["task"],
"system": prompt["system"],
"user": prompt["user"],
"response": clean(content),
"metadata": {"style": prompt["style"]}
}
except Exception:
await asyncio.sleep(0.5)
return None
async def main():
count = int(sys.argv[1]) if len(sys.argv) > 1 else 500
semaphore = asyncio.Semaphore(10)
prompts = [build_prompt(random.choice(STYLES)) for _ in range(count)]
results = []
async with aiohttp.ClientSession() as session:
tasks = [call_one(session, p, semaphore) for p in prompts]
for i, coro in enumerate(asyncio.as_completed(tasks)):
result = await coro
if result:
results.append(result)
if (i + 1) % 50 == 0 or (i + 1) == len(tasks):
print(f"[{i+1}/{len(tasks)}] valid: {len(results)}")
with open(EXTRA_FILE, "w", encoding="utf-8") as f:
for row in results:
f.write(json.dumps(row, ensure_ascii=False) + "\n")
print(f"Generated {len(results)} extra mentor rows -> {EXTRA_FILE}")
if __name__ == "__main__":
asyncio.run(main())
|