alphaear-reporter — Skillopedia

AlphaEar Reporter Skill Overview This skill provides a structured workflow for generating professional financial reports. It includes planning, writing, editing, and creating visual aids (charts). Capabilities Capabilities 1. Generate Structured Reports (Agentic Workflow) YOU (the Agent) are the Report Generator. Use the prompts in to progressively build the report. Workflow: 1. Cluster Signals : Read input signals and use the Cluster Signals Prompt to group them. 2. Write Sections : For each cluster, use the Write Section Prompt to generate analysis. 3. Assemble : Use the Final Assembly Prom…

, '', text)\n \n # 2. 寻找第一个 JSON 起始符 { 或 [\n start_brace = text.find('{')\n start_bracket = text.find('[')\n \n if start_brace == -1 and start_bracket == -1:\n return None\n \n start_idx = start_brace if (start_bracket == -1 or (start_brace != -1 and start_brace \u003c start_bracket)) else start_bracket\n \n # 2.5 预处理：修复一些极其常见的 LLM 错误\n potential_json = text[start_idx:].strip()\n \n # remove comments safely\n potential_json = _strip_comments(potential_json)\n \n # b. 修复缺失开头引号的键: nodes\": [ -> \"nodes\": [\n # 匹配模式: (空白或换行) 单词紧跟引号和冒号\n potential_json = re.sub(r'([\\{\\,]\\s*)([a-zA-Z_]\\w*)\\\"\\s*:', r'\\1\"\\2\":', potential_json)\n \n # c. 修复缺失末尾引号的键: \"nodes: [ -> \"nodes\": [\n potential_json = re.sub(r'([\\{\\,]\\s*)\\\"([a-zA-Z_]\\w*)\\s*:', r'\\1\"\\2\":', potential_json)\n\n # d. 修复完全缺失引号的键: nodes: [ -> \"nodes\": [\n # 注意避免匹配到像 http:// 这种内容，所以限定在 { 或 , 之后\n potential_json = re.sub(r'([\\{\\,]\\s*)([a-zA-Z_]\\w*)\\s*:', r'\\1\"\\2\":', potential_json)\n \n # 3. 使用 raw_decode 尝试解析\n decoder = json.JSONDecoder()\n \n # 首先尝试直接解析（不做任何预处理）\n try:\n obj = json.loads(potential_json)\n return obj\n except json.JSONDecodeError:\n pass\n \n # 简单预处理：移除对象/列表末位多余逗号\n processed_json = re.sub(r',\\s*([\\]}])', r'\\1', potential_json)\n \n try:\n obj, end_pos = decoder.raw_decode(processed_json)\n return obj\n except json.JSONDecodeError:\n pass\n \n # e. 修复未终止的字符串字面量问题：移除值中的实际换行符\n # LLM 可能在字符串值中生成包含真实 newline 的内容，导致 JSON 非法\n def fix_multiline_strings(s):\n # 简单策略：将字符串值内的换行替换为空格\n lines = s.split('\\n')\n result = []\n in_string = False\n for line in lines:\n # 计算未转义的引号数\n quote_count = line.count('\"') - line.count('\\\\\"')\n if in_string:\n result[-1] += ' ' + line.strip()\n else:\n result.append(line)\n \n if quote_count % 2 == 1:\n in_string = not in_string\n return '\\n'.join(result)\n \n fixed_json = fix_multiline_strings(processed_json)\n \n try:\n obj, end_pos = decoder.raw_decode(fixed_json)\n return obj\n except json.JSONDecodeError:\n try:\n # 4. 尝试处理单引号问题 (JSON 规范要求双引号，但 LLM 常输出单引号)\n # 这是一个简单的替换技巧，仅针对像 {'key': 'value'} 这样的结构\n # 注意：这可能会破坏包含单引号的字符串值，所以作为较后的回退\n fix_quotes = re.sub(r\"'(.*?)':\", r'\"\\1\":', processed_json) # 修复键\n fix_quotes = re.sub(r\":\\s*'(.*?)'\", r': \"\\1\"', fix_quotes) # 修复简单值\n obj, end_pos = decoder.raw_decode(fix_quotes)\n return obj\n except (json.JSONDecodeError, TypeError):\n try:\n # 5. 使用 ast.literal_eval 作为终极回退 (处理 Python 字典格式)\n # 提取第一个匹配的括号对内容\n # 寻找匹配的 { }\n stack = []\n for i, char in enumerate(potential_json):\n if char == '{': stack.append('{')\n elif char == '}':\n if stack: stack.pop()\n if not stack:\n content = potential_json[:i+1]\n return ast.literal_eval(content)\n except (ValueError, SyntaxError, MemoryError) as e:\n logger.warning(f\"All JSON extraction attempts failed: {e}\")\n except Exception as e:\n logger.error(f\"Unexpected error during JSON extraction: {e}\")\n \n return None\n","content_type":"text/x-python; charset=utf-8","language":"python","size":6300,"content_sha256":"2ec52de0a706f64b3466f490dc3f12ad8241f20fa316e858f7e3393afb3062d9"},{"filename":"scripts/utils/llm/capability.py","content":"import os\nfrom typing import Optional, List, Dict, Any\nfrom agno.agent import Agent\nfrom agno.models.base import Model\nfrom loguru import logger\nfrom ..llm.factory import get_model\n\ndef test_tool_call_support(model: Model) -> bool:\n \"\"\"\n 测试模型是否支持原生的 Tool Call (Function Calling)。\n 通过尝试执行一个简单的加法工具来验证。\n \"\"\"\n def get_current_weather(location: str):\n \"\"\"获取指定地点的天气\"\"\"\n return f\"{location} 的天气是晴天，25度。\"\n\n test_agent = Agent(\n model=model,\n tools=[get_current_weather],\n instructions=\"请调用工具查询北京的天气，并直接返回工具的输出结果。\"\n )\n\n try:\n # 运行一个简单的任务，观察是否触发了 tool_call\n response = test_agent.run(\"北京天气怎么样？\")\n \n # 检查 response 中是否包含 tool_calls\n # Agno 的 RunResponse 对象通常包含 messages，我们可以检查最后几条消息\n has_tool_call = False\n for msg in response.messages:\n if hasattr(msg, 'tool_calls') and msg.tool_calls:\n has_tool_call = True\n break\n \n if has_tool_call:\n logger.info(f\"✅ Model {model.id} supports native tool calling.\")\n return True\n else:\n # 如果没有 tool_calls 但返回了正确答案，可能是模型通过纯文本模拟了工具调用（ReAct）\n # 或者根本没用工具。对于原生支持的判断，我们坚持要求有 tool_calls 结构。\n logger.warning(f\"⚠️ Model {model.id} did NOT use native tool calling structure.\")\n return False\n \n except Exception as e:\n logger.error(f\"❌ Error testing tool call for {model.id}: {e}\")\n return False\n\nclass ModelCapabilityRegistry:\n \"\"\"\n 模型能力注册表，用于缓存和管理不同模型的能力测试结果。\n \"\"\"\n _cache = {}\n\n @classmethod\n def get_capabilities(cls, provider: str, model_id: str, **kwargs) -> Dict[str, bool]:\n key = f\"{provider}:{model_id}\"\n if key not in cls._cache:\n logger.info(f\"🔍 Testing capabilities for {key}...\")\n model = get_model(provider, model_id, **kwargs)\n supports_tool_call = test_tool_call_support(model)\n cls._cache[key] = {\n \"supports_tool_call\": supports_tool_call\n }\n return cls._cache[key]\n\nif __name__ == \"__main__\":\n # 简单测试脚本\n from dotenv import load_dotenv\n load_dotenv()\n \n # 测试当前配置的模型\n p = os.getenv(\"LLM_PROVIDER\", \"ust\")\n m = os.getenv(\"LLM_MODEL\", \"Qwen\")\n \n print(f\"Testing {p}/{m}...\")\n res = ModelCapabilityRegistry.get_capabilities(p, m)\n print(f\"Result: {res}\")\n","content_type":"text/x-python; charset=utf-8","language":"python","size":2863,"content_sha256":"40b215695e2b7519224773ec4b280d115770670d1ddfc37e3eea2d23e085e171"},{"filename":"scripts/utils/llm/factory.py","content":"import os\nfrom agno.models.openai import OpenAIChat\nfrom agno.models.ollama import Ollama\nfrom agno.models.dashscope import DashScope\nfrom agno.models.deepseek import DeepSeek\nfrom agno.models.openrouter import OpenRouter\n\ndef get_model(model_provider: str, model_id: str, **kwargs):\n \"\"\"\n Factory to get the appropriate LLM model.\n \n Args:\n model_provider: \"openai\", \"ollama\", \"deepseek\"\n model_id: The specific model ID (e.g., \"gpt-4o\", \"llama3\", \"deepseek-chat\")\n **kwargs: Additional arguments for the model constructor\n \"\"\"\n if model_provider == \"openai\":\n return OpenAIChat(id=model_id, **kwargs)\n \n elif model_provider == \"ollama\":\n return Ollama(id=model_id, **kwargs)\n \n elif model_provider == \"deepseek\":\n # DeepSeek is OpenAI compatible\n api_key = os.getenv(\"DEEPSEEK_API_KEY\")\n if not api_key:\n print(\"Warning: DEEPSEEK_API_KEY not set.\")\n \n return DeepSeek(\n id=model_id,\n api_key=api_key,\n **kwargs\n )\n elif model_provider == \"dashscope\":\n api_key = os.getenv(\"DASHSCOPE_API_KEY\")\n if not api_key:\n print(\"Warning: DASHSCOPE_API_KEY not set.\")\n \n return DashScope(\n id=model_id,\n base_url=\"https://dashscope.aliyuncs.com/compatible-mode/v1\",\n api_key=api_key,\n **kwargs\n )\n elif model_provider == 'openrouter':\n api_key = os.getenv(\"OPENROUTER_API_KEY\")\n if not api_key:\n print('Warning: OPENROUTER_API_KEY not set.')\n \n return OpenRouter(\n id=model_id,\n api_key=api_key,\n **kwargs\n )\n\n elif model_provider == 'zai':\n api_key = os.getenv(\"ZAI_KEY_API\")\n if not api_key:\n print('Warning: ZAI_KEY_API not set.')\n\n # role_map to ensure compatibility.\n default_role_map = {\n \"system\": \"system\",\n \"user\": \"user\",\n \"assistant\": \"assistant\",\n \"tool\": \"tool\",\n \"model\": \"assistant\",\n }\n\n # Allow callers to override role_map via kwargs, otherwise use default\n role_map = kwargs.pop(\"role_map\", default_role_map)\n \n return OpenAIChat(\n id=model_id,\n base_url=\"https://api.z.ai/api/paas/v4\",\n api_key=api_key,\n timeout=60,\n role_map=role_map,\n extra_body={\"enable_thinking\": False}, # TODO: one more setting for thinking\n **kwargs\n )\n \n elif model_provider == 'ust':\n api_key = os.getenv(\"UST_KEY_API\")\n if not api_key:\n print('Warning: UST_KEY_API not set.')\n \n # Some UST-compatible endpoints expect the standard OpenAI role names\n # (e.g. \"system\", \"user\", \"assistant\") rather than Agno's default\n # mapping which maps \"system\" -> \"developer\". Provide an explicit\n # role_map to ensure compatibility.\n default_role_map = {\n \"system\": \"system\",\n \"user\": \"user\",\n \"assistant\": \"assistant\",\n \"tool\": \"tool\",\n \"model\": \"assistant\",\n }\n\n # Allow callers to override role_map via kwargs, otherwise use default\n role_map = kwargs.pop(\"role_map\", default_role_map)\n\n return OpenAIChat(\n id=model_id,\n api_key=api_key,\n base_url=os.getenv(\"UST_URL\"),\n role_map=role_map,\n extra_body={\"enable_thinking\": False}, # TODO: one more setting for thinking\n **kwargs\n )\n \n else:\n raise ValueError(f\"Unknown model provider: {model_provider}\")\n\n","content_type":"text/x-python; charset=utf-8","language":"python","size":3717,"content_sha256":"85065ea8a843075f86a2532d455e64dca2207b657b758196d89cff9fad9e67b0"},{"filename":"scripts/utils/llm/router.py","content":"import os\nfrom typing import Optional, List, Dict, Any, Union\nfrom agno.models.base import Model\nfrom loguru import logger\nfrom dotenv import load_dotenv\nfrom ..llm.factory import get_model\nfrom ..llm.capability import ModelCapabilityRegistry\n\n# 确保在初始化前加载环境变量\nload_dotenv()\n\nclass ModelRouter:\n \"\"\"\n 模型路由管理器\n \n 功能：\n 1. 管理“推理/写作模型” (Reasoning Model) 和“工具调用模型” (Tool Model)。\n 2. 根据任务需求自动选择合适的模型。\n \"\"\"\n \n def __init__(self):\n # 默认从环境变量读取\n self.reasoning_provider = os.getenv(\"REASONING_MODEL_PROVIDER\", os.getenv(\"LLM_PROVIDER\", \"openai\"))\n self.reasoning_id = os.getenv(\"REASONING_MODEL_ID\", os.getenv(\"LLM_MODEL\", \"gpt-4o\"))\n self.reasoning_host = os.getenv(\"REASONING_MODEL_HOST\", os.getenv(\"LLM_HOST\"))\n \n self.tool_provider = os.getenv(\"TOOL_MODEL_PROVIDER\", self.reasoning_provider)\n self.tool_id = os.getenv(\"TOOL_MODEL_ID\", self.reasoning_id)\n self.tool_host = os.getenv(\"TOOL_MODEL_HOST\", self.reasoning_host)\n \n self._reasoning_model = None\n self._tool_model = None\n \n logger.info(f\"🤖 ModelRouter initialized: Reasoning={self.reasoning_id} ({self.reasoning_host or 'default'}), Tool={self.tool_id} ({self.tool_host or 'default'})\")\n\n def get_reasoning_model(self, **kwargs) -> Model:\n if not self._reasoning_model:\n # 优先使用路由配置的 host\n if self.reasoning_host and \"host\" not in kwargs:\n kwargs[\"host\"] = self.reasoning_host\n self._reasoning_model = get_model(self.reasoning_provider, self.reasoning_id, **kwargs)\n return self._reasoning_model\n\n def get_tool_model(self, **kwargs) -> Model:\n if not self._tool_model:\n # 优先使用路由配置的 host\n if self.tool_host and \"host\" not in kwargs:\n kwargs[\"host\"] = self.tool_host\n \n # 检查 tool_model 是否真的支持 tool call\n caps = ModelCapabilityRegistry.get_capabilities(self.tool_provider, self.tool_id, **kwargs)\n if not caps[\"supports_tool_call\"]:\n logger.warning(f\"⚠️ Configured tool model {self.tool_id} might not support native tool calls! Consider using ReAct mode or a different model.\")\n \n self._tool_model = get_model(self.tool_provider, self.tool_id, **kwargs)\n return self._tool_model\n\n def get_model_for_agent(self, has_tools: bool = False, **kwargs) -> Model:\n \"\"\"\n 根据 Agent 是否包含工具来返回合适的模型。\n \"\"\"\n if has_tools:\n return self.get_tool_model(**kwargs)\n return self.get_reasoning_model(**kwargs)\n\n# 全局单例\nrouter = ModelRouter()\n","content_type":"text/x-python; charset=utf-8","language":"python","size":2875,"content_sha256":"0a8dbfa1a9d8d2a5b62ca77bc5a78bbbb1e53706be9cf18a7290e7608649927a"},{"filename":"scripts/utils/logging_setup.py","content":"import os\nimport sys\nfrom datetime import datetime\nfrom typing import Optional\n\nfrom loguru import logger\n\n\ndef setup_file_logging(\n run_id: str,\n log_dir: str = \"logs\",\n level: str = \"INFO\",\n retention: str = \"10 days\",\n rotation: str = \"20 MB\",\n) -> str:\n \"\"\"Configure Loguru to log to stderr + a per-run file.\n\n Returns the log file path.\n \"\"\"\n os.makedirs(log_dir, exist_ok=True)\n\n # Remove default handler to avoid duplicate logs.\n logger.remove()\n\n # Console\n logger.add(sys.stderr, level=level, backtrace=False, diagnose=False)\n\n # File (safe for multi-thread via enqueue)\n log_path = os.path.join(log_dir, f\"signalflux_{run_id}.log\")\n logger.add(\n log_path,\n level=level,\n rotation=rotation,\n retention=retention,\n enqueue=True,\n backtrace=True,\n diagnose=False,\n encoding=\"utf-8\",\n )\n return log_path\n\n\ndef make_run_id(prefix: Optional[str] = None) -> str:\n ts = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n return f\"{prefix}_{ts}\" if prefix else ts\n","content_type":"text/x-python; charset=utf-8","language":"python","size":1072,"content_sha256":"2b0e1a87da39920ed0958e5f9915cf481dccc8016874c91acaf2a0403916eadc"},{"filename":"scripts/utils/news_tools.py","content":"import requests\nfrom requests.exceptions import RequestException, Timeout\nimport json\nimport time\nfrom datetime import datetime\nfrom typing import List, Dict, Optional\nfrom loguru import logger\nfrom .database_manager import DatabaseManager\nfrom .content_extractor import ContentExtractor\n\nclass NewsNowTools:\n \"\"\"热点新闻获取工具 - 接入 NewsNow API 与 Jina 内容提取\"\"\"\n \n BASE_URL = \"https://newsnow.busiyi.world\"\n SOURCES = {\n # 金融类\n \"cls\": \"财联社\",\n \"wallstreetcn\": \"华尔街见闻\",\n \"xueqiu\": \"雪球热榜\",\n # 综合/社交\n \"weibo\": \"微博热搜\",\n \"zhihu\": \"知乎热榜\",\n \"baidu\": \"百度热搜\",\n \"toutiao\": \"今日头条\",\n \"douyin\": \"抖音热榜\",\n \"thepaper\": \"澎湃新闻\",\n # 科技类\n \"36kr\": \"36氪\",\n \"ithome\": \"IT之家\",\n \"v2ex\": \"V2EX\",\n \"juejin\": \"掘金\",\n \"hackernews\": \"Hacker News\",\n }\n\n\n def __init__(self, db: DatabaseManager):\n self.db = db\n self.user_agent = (\n \"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) \"\n \"AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36\"\n )\n self.extractor = ContentExtractor()\n # Simple in-memory cache: source_id -> {\"time\": timestamp, \"data\": []}\n self._cache = {}\n\n def fetch_hot_news(self, source_id: str, count: int = 15, fetch_content: bool = False) -> List[Dict]:\n \"\"\"\n 从指定新闻源获取热点新闻列表（支持5分钟缓存）。\n \"\"\"\n # 1. Check cache validity (5 minutes)\n cache_key = f\"{source_id}_{count}\"\n cached = self._cache.get(cache_key)\n now = time.time()\n \n if cached and (now - cached[\"time\"] \u003c 300):\n logger.info(f\"⚡ Using cached news for {source_id} (Age: {int(now - cached['time'])}s)\")\n return cached[\"data\"]\n\n try:\n url = f\"{self.BASE_URL}/api/s?id={source_id}\"\n response = requests.get(url, headers={\"User-Agent\": self.user_agent}, timeout=30)\n if response.status_code == 200:\n data = response.json()\n items = data.get(\"items\", [])[:count]\n processed_items = []\n for i, item in enumerate(items, 1):\n item_url = item.get(\"url\", \"\")\n content = \"\"\n if fetch_content and item_url:\n content = self.extractor.extract_with_jina(item_url) or \"\"\n \n processed_items.append({\n \"id\": item.get(\"id\") or f\"{source_id}_{int(time.time())}_{i}\",\n \"source\": source_id,\n \"rank\": i,\n \"title\": item.get(\"title\", \"\"),\n \"url\": item_url,\n \"content\": content,\n \"publish_time\": item.get(\"publish_time\"),\n \"meta_data\": item.get(\"extra\", {})\n })\n \n # Update Cache\n self._cache[cache_key] = {\"time\": now, \"data\": processed_items}\n logger.info(f\"✅ Fetched and cached news for {source_id}\")\n \n self.db.save_daily_news(processed_items)\n return processed_items\n else:\n logger.error(f\"NewsNow API Error: {response.status_code}\")\n # Fallback to stale cache if available\n if cached:\n logger.warning(f\"⚠️ API failed, using stale cache for {source_id}\")\n return cached[\"data\"]\n return []\n except Timeout:\n logger.error(f\"Timeout fetching hot news from {source_id}\")\n if cached:\n logger.warning(f\"⚠️ Timeout, using stale cache for {source_id}\")\n return cached[\"data\"]\n return []\n except RequestException as e:\n logger.error(f\"Network error fetching hot news from {source_id}: {e}\")\n if cached:\n logger.warning(f\"⚠️ Network check failed, using stale cache for {source_id}\")\n return cached[\"data\"]\n return []\n except json.JSONDecodeError:\n logger.error(f\"Failed to parse JSON response from NewsNow for {source_id}\")\n return []\n except Exception as e:\n logger.error(f\"Unexpected error fetching hot news from {source_id}: {e}\")\n return []\n\n def fetch_news_content(self, url: str) -> Optional[str]:\n \"\"\"\n 使用 Jina Reader 抓取指定 URL 的网页正文内容。\n \n Args:\n url: 需要抓取内容的完整网页 URL，必须以 http:// 或 https:// 开头。\n \n Returns:\n 提取的网页正文内容 (Markdown 格式)，如果失败则返回 None。\n \"\"\"\n return self.extractor.extract_with_jina(url)\n\n def get_unified_trends(self, sources: Optional[List[str]] = None) -> str:\n \"\"\"\n 获取多平台综合热点报告，自动聚合多个新闻源的热门内容。\n \n Args:\n sources: 要扫描的新闻源列表。可选值按类别:\n **金融类**: \"cls\", \"wallstreetcn\", \"xueqiu\"\n **综合类**: \"weibo\", \"zhihu\", \"baidu\", \"toutiao\", \"douyin\", \"thepaper\"\n **科技类**: \"36kr\", \"ithome\", \"v2ex\", \"juejin\", \"hackernews\"\n \n Returns:\n 格式化的 Markdown 热点汇总报告，包含各平台 Top 10 热点标题和链接。\n \"\"\"\n sources = sources or [\"weibo\", \"zhihu\", \"wallstreetcn\"]\n all_news = []\n for src in sources:\n all_news.extend(self.fetch_hot_news(src))\n time.sleep(0.2)\n \n if not all_news:\n return \"❌ 未能获取到热点数据\"\n \n report = f\"# 实时全网热点汇总 ({datetime.now().strftime('%Y-%m-%d %H:%M')})\\n\\n\"\n for src in sources:\n\n src_name = self.SOURCES.get(src, src)\n report += f\"### 🔥 {src_name}\\n\"\n src_news = [n for n in all_news if n['source'] == src]\n for n in src_news[:10]:\n report += f\"- {n['title']} ([链接]({n['url']}))\\n\"\n report += \"\\n\"\n \n return report\n\n\nclass PolymarketTools:\n \"\"\"Polymarket 预测市场数据工具 - 获取热门预测市场反映公众情绪和预期\"\"\"\n \n BASE_URL = \"https://gamma-api.polymarket.com\"\n \n def __init__(self, db: DatabaseManager):\n self.db = db\n self.user_agent = \"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36\"\n \n def get_active_markets(self, limit: int = 20) -> List[Dict]:\n \"\"\"\n 获取活跃的预测市场，用于分析公众情绪和预期。\n \n 预测市场数据可以反映:\n - 公众对重大事件的预期概率\n - 市场情绪和风险偏好\n - 热门话题的关注度\n \n Args:\n limit: 获取的市场数量，默认 20 个。\n \n Returns:\n 包含预测市场信息的列表，每个市场包含:\n - question: 预测问题\n - outcomes: 可能的结果\n - outcomePrices: 各结果的概率价格\n - volume: 交易量\n \"\"\"\n try:\n response = requests.get(\n f\"{self.BASE_URL}/markets\",\n params={\"active\": \"true\", \"closed\": \"false\", \"limit\": limit},\n headers={\"User-Agent\": self.user_agent, \"Accept\": \"application/json\"},\n timeout=30\n )\n \n if response.status_code == 200:\n markets = response.json()\n result = []\n for m in markets:\n result.append({\n \"id\": m.get(\"id\"),\n \"question\": m.get(\"question\"),\n \"slug\": m.get(\"slug\"),\n \"outcomes\": m.get(\"outcomes\"),\n \"outcomePrices\": m.get(\"outcomePrices\"),\n \"volume\": m.get(\"volume\"),\n \"liquidity\": m.get(\"liquidity\"),\n })\n logger.info(f\"✅ 获取 {len(result)} 个预测市场\")\n return result\n else:\n logger.warning(f\"⚠️ Polymarket API 返回 {response.status_code}\")\n return []\n except Timeout:\n logger.error(\"Timeout fetching Polymarket markets\")\n return []\n except RequestException as e:\n logger.error(f\"Network error fetching Polymarket markets: {e}\")\n return []\n except json.JSONDecodeError:\n logger.error(\"Failed to parse JSON response from Polymarket\")\n return []\n except Exception as e:\n logger.error(f\"Unexpected error fetching Polymarket markets: {e}\")\n return []\n \n def get_market_summary(self, limit: int = 10) -> str:\n \"\"\"\n 获取预测市场摘要报告，用于了解当前热门话题和公众预期。\n \n Args:\n limit: 获取的市场数量\n \n Returns:\n 格式化的预测市场报告\n \"\"\"\n markets = self.get_active_markets(limit)\n if not markets:\n return \"❌ 无法获取 Polymarket 数据\"\n \n report = f\"# 🔮 Polymarket 热门预测 ({datetime.now().strftime('%Y-%m-%d %H:%M')})\\n\\n\"\n for i, m in enumerate(markets, 1):\n question = m.get(\"question\", \"Unknown\")\n prices = m.get(\"outcomePrices\", [])\n volume = m.get(\"volume\", 0)\n \n report += f\"**{i}. {question}**\\n\"\n if prices:\n report += f\" 概率: {prices}\\n\"\n if volume:\n report += f\" 交易量: ${float(volume):,.0f}\\n\"\n report += \"\\n\"\n \n return report\n","content_type":"text/x-python; charset=utf-8","language":"python","size":10109,"content_sha256":"85e5b0a3a57f1cb5ed9105404fbb2de527639a4bda9d433af42a3a48ff1f71eb"},{"filename":"scripts/utils/predictor/evaluation.py","content":"import os\nimport sys\nimport torch\nimport pandas as pd\nimport numpy as np\nimport glob\nfrom loguru import logger\nfrom datetime import datetime, timedelta\n\n# Setup paths\nKRONOS_DIR = os.path.dirname(os.path.abspath(__file__))\nSRC_DIR = os.path.dirname(os.path.dirname(KRONOS_DIR))\nif SRC_DIR not in sys.path:\n sys.path.insert(0, SRC_DIR)\n\nfrom ..kronos.auto_synthesis_training import AutoSynthesisTrainer\nfrom ..kronos.model import KronosPredictor\nfrom ..visualizer import VisualizerTools\nfrom ..schema.models import ForecastResult, KLinePoint\n\nclass NewsModelEvaluator:\n def __init__(self, model_path=None):\n self.trainer = AutoSynthesisTrainer()\n self.device = self.trainer.device\n \n if model_path is None:\n # Try to find the latest model in exports/models\n model_files = glob.glob(os.path.join(SRC_DIR, \"exports/models/*.pt\"))\n if not model_files:\n logger.warning(\"⚠️ No trained models found in exports/models/. Using base model (zero-init proj).\")\n else:\n model_path = max(model_files, key=os.path.getctime)\n \n if model_path:\n self.load_weights(model_path)\n\n def load_weights(self, path):\n logger.info(f\"🔄 Loading model weights from {path}...\")\n checkpoint = torch.load(path, map_location=self.device)\n self.trainer.model.news_proj.load_state_dict(checkpoint['news_proj_state_dict'])\n logger.success(\"✅ News projection layer loaded.\")\n\n def evaluate_range(self, start_idx=100, end_idx=200, pred_len=5):\n # 1. Fetch Tickers\n res = self.trainer.db.execute_query(\"SELECT code FROM stock_list\")\n all_tickers = [row['code'] for row in res]\n test_tickers = all_tickers[start_idx:end_idx]\n \n if not test_tickers:\n logger.error(f\"No tickers found in range {start_idx}-{end_idx}\")\n return\n\n logger.info(f\"🚀 Evaluating News Model on stocks {start_idx} to {end_idx}...\")\n \n # 2. Discover Shocks\n shocks = self.trainer.discover_shocks(test_tickers, pred_len=pred_len)\n \n # 3. Associate News & Predict\n self.trainer.model.eval()\n predictor = KronosPredictor(self.trainer.model, self.trainer.tokenizer, device=self.device)\n \n save_dir = os.path.join(SRC_DIR, \"exports/evaluation_results\")\n os.makedirs(save_dir, exist_ok=True)\n\n count = 0\n for shock in shocks:\n summary = self.trainer.find_reason_and_verify(shock)\n if not summary:\n continue\n \n logger.info(f\"📈 Testing shock: {shock['ticker']} on {shock['date']}\")\n \n # Embedding news\n news_emb = self.trainer.embedder.encode(summary)\n \n # Prediction\n h = shock['history']\n t = shock['target']\n actuals = t['close'].values[:pred_len]\n \n x_ts = pd.to_datetime(h['date'])\n future_dates = pd.date_range(start=x_ts.iloc[-1] + timedelta(days=1), periods=pred_len, freq='B')\n y_ts = pd.Series(future_dates)\n \n # A. Base Prediction (No news)\n p_base = predictor.predict(h, x_ts, y_ts, pred_len=pred_len, news_emb=None, verbose=False)\n \n # B. News-Aware Prediction\n p_news = predictor.predict(h, x_ts, y_ts, pred_len=pred_len, news_emb=news_emb, verbose=False)\n \n # Calculate Improvement\n b_preds = p_base['close'].values[:len(actuals)]\n n_preds = p_news['close'].values[:len(actuals)]\n b_mae = np.mean(np.abs(b_preds - actuals))\n n_mae = np.mean(np.abs(n_preds - actuals))\n improvement = (b_mae - n_mae) / (b_mae + 1e-6) * 100\n\n # C. Visualize\n try:\n def to_kp_list(preds_df):\n points = []\n for idx, row in preds_df.iterrows():\n points.append(KLinePoint(\n date=str(idx)[:10], open=row['open'], high=row['high'],\n low=row['low'], close=row['close'], volume=row.get('volume', 0)\n ))\n return points\n\n forecast_obj = ForecastResult(\n ticker=shock['ticker'],\n base_forecast=to_kp_list(p_base),\n adjusted_forecast=to_kp_list(p_news),\n rationale=summary\n )\n\n chart = VisualizerTools.generate_stock_chart(\n df=h, ticker=shock['ticker'],\n title=f\"Test Eval: {shock['ticker']} ({shock['date']}) Imp: {improvement:.1f}%\",\n forecast=forecast_obj,\n ground_truth=t[['date', 'open', 'high', 'low', 'close', 'volume']]\n )\n \n safe_date = shock['date'].replace(\"-\", \"\")\n filename = f\"test_{shock['ticker']}_{safe_date}.html\"\n VisualizerTools.render_chart_to_file(chart, os.path.join(save_dir, filename))\n \n logger.success(f\"📊 Result for {shock['ticker']} saved. Base MAE: {b_mae:.4f}, News MAE: {n_mae:.4f}\")\n count += 1\n except Exception as e:\n logger.error(f\"Visualization failed: {e}\")\n\n logger.info(f\"🏁 Finished evaluation. {count} cases visualized in {save_dir}\")\n\nif __name__ == \"__main__\":\n # If you have a specific model, pass the path here. Otherwise it picks the latest.\n evaluator = NewsModelEvaluator()\n evaluator.evaluate_range(start_idx=100, end_idx=200, pred_len=1)\n","content_type":"text/x-python; charset=utf-8","language":"python","size":5723,"content_sha256":"c00bbac8132e737bf81917ad79dbc1dfe902f66c9d8b1813649d4bad6cbec78d"},{"filename":"scripts/utils/predictor/kline_generate.py","content":"# Ref: https://github.com/shiyu-coder/Kronos\n\nfrom model import Kronos, KronosTokenizer, KronosPredictor\nimport pandas as pd\nimport sqlite3\nimport torch\nimport matplotlib.pyplot as plt\nimport matplotlib.gridspec as gridspec\nfrom pandas.tseries.offsets import BusinessDay\nimport numpy as np\n\ndef get_device():\n device = \"cuda\" if torch.cuda.is_available() else \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n print(f\"Using device: {device}\")\n return device\n\ndef load_predictor():\n tokenizer = KronosTokenizer.from_pretrained(\"NeoQuasar/Kronos-Tokenizer-base\")\n model = Kronos.from_pretrained(\"NeoQuasar/Kronos-base\")\n device = get_device()\n tokenizer = tokenizer.to(device)\n model = model.to(device)\n return KronosPredictor(model, tokenizer, device=device, max_context=512)\n\ndef load_data(ticker=\"002111\", db_path=\"AlphaEar/data/signal_flux.db\"):\n with sqlite3.connect(db_path) as conn:\n df = pd.read_sql_query(f\"SELECT * FROM stock_prices WHERE ticker = '{ticker}'\", conn)\n df['date'] = pd.to_datetime(df['date'])\n df = df.sort_values('date').reset_index(drop=True)\n return df\n\ndef plot_kline_matplotlib(ax, ax_vol, dates, df, label_suffix=\"\", color_up='#ef4444', color_down='#22c55e', alpha=1.0, is_prediction=False):\n \"\"\"\n 绘制 K 线图和成交量\n \"\"\"\n # X axis mapping to integers for consistent spacing\n x = np.arange(len(dates))\n \n # K-line data\n opens = df['open'].values\n closes = df['close'].values\n highs = df['high'].values\n lows = df['low'].values\n volumes = df['volume'].values\n \n # Width of the candlestick\n width = 0.6\n \n for i in range(len(x)):\n color = color_up if closes[i] >= opens[i] else color_down\n linestyle = '--' if is_prediction else '-'\n \n # Wick\n ax.vlines(x[i], lows[i], highs[i], color=color, linewidth=1, alpha=alpha, linestyle=linestyle)\n \n # Body\n rect_bottom = min(opens[i], closes[i])\n rect_height = abs(opens[i] - closes[i])\n if rect_height == 0: rect_height = 0.001 # Visual hair\n \n ax.add_patch(plt.Rectangle((x[i] - width/2, rect_bottom), width, rect_height, \n edgecolor=color, facecolor=color if not is_prediction else 'none', \n alpha=alpha, linewidth=1, linestyle=linestyle))\n \n # Volume\n ax_vol.bar(x[i], volumes[i], color=color, alpha=alpha * 0.5, width=width)\n\ndef render_comparison_chart(history_df, actual_df, pred_df, title):\n \"\"\"\n 渲染组合图：历史 K 线 + 真值 K 线 + 预测 K 线\n \"\"\"\n # Combine all dates for X axis\n all_dates = pd.concat([history_df['date'], actual_df['date'] if actual_df is not None else pred_df.index.to_series()]).unique()\n all_dates = sorted(all_dates)\n date_to_idx = {date: i for i, date in enumerate(all_dates)}\n \n fig = plt.figure(figsize=(14, 8), facecolor='white')\n gs = gridspec.GridSpec(2, 1, height_ratios=[3, 1], hspace=0.1)\n ax_main = fig.add_subplot(gs[0])\n ax_vol = fig.add_subplot(gs[1], sharex=ax_main)\n \n # 1. Plot History\n hist_indices = [date_to_idx[d] for d in history_df['date']]\n # We use a custom x for plotting to ensure continuity\n plot_kline_matplotlib(ax_main, ax_vol, history_df['date'], history_df, alpha=0.8)\n \n offset = len(history_df)\n \n # 2. Plot Actual if exists\n if actual_df is not None:\n # Shift indices\n actual_x = np.arange(len(actual_df)) + offset\n # Plotting manually to handle offset\n for i in range(len(actual_df)):\n idx = actual_x[i]\n row = actual_df.iloc[i]\n color = '#ef4444' if row['close'] >= row['open'] else '#22c55e'\n ax_main.vlines(idx, row['low'], row['high'], color=color, linewidth=1, alpha=0.9)\n ax_main.add_patch(plt.Rectangle((idx - 0.3, min(row['open'], row['close'])), 0.6, abs(row['open']-row['close']), \n edgecolor=color, facecolor=color, alpha=0.9))\n ax_vol.bar(idx, row['volume'], color=color, alpha=0.4)\n \n # 3. Plot Prediction\n pred_x = np.arange(len(pred_df)) + offset\n for i in range(len(pred_df)):\n idx = pred_x[i]\n row = pred_df.iloc[i]\n color = '#ff8c00' # Orange for prediction to distinguish\n ax_main.vlines(idx, row['low'], row['high'], color=color, linewidth=1.5, linestyle='--')\n ax_main.add_patch(plt.Rectangle((idx - 0.3, min(row['open'], row['close'])), 0.6, abs(row['open']-row['close']), \n edgecolor=color, facecolor='none', linewidth=1.5, linestyle='--'))\n # Plot secondary prediction line for close\n if i == 0:\n # Connect to history\n ax_main.plot([offset-1, idx], [history_df['close'].iloc[-1], row['close']], color=color, linestyle='--', alpha=0.6)\n elif i > 0:\n ax_main.plot([idx-1, idx], [pred_df['close'].iloc[i-1], row['close']], color=color, linestyle='--', alpha=0.6)\n\n # Styling\n ax_main.set_title(title, fontsize=14, fontweight='bold')\n ax_main.grid(True, linestyle=':', alpha=0.6)\n ax_vol.grid(True, linestyle=':', alpha=0.6)\n ax_vol.set_ylabel('Volume')\n ax_main.set_ylabel('Price')\n \n # Set X ticks\n step = max(1, len(all_dates) // 10)\n ax_vol.set_xticks(np.arange(0, len(all_dates), step))\n ax_vol.set_xticklabels([all_dates[i].strftime('%Y-%m-%d') for i in range(0, len(all_dates), step)], rotation=45)\n \n plt.tight_layout()\n plt.show()\n plt.close()\n\ndef run_backtest(df, predictor, lookback, pred_len, start_index=0):\n total_len = len(df)\n history_start = start_index\n history_end = start_index + lookback \n pred_start = history_end\n \n available_pred_len = total_len - pred_start\n if available_pred_len \u003c= 0: return\n actual_pred_len = min(pred_len, available_pred_len)\n pred_end = pred_start + actual_pred_len\n \n x_df = df.iloc[history_start : history_end].copy()\n y_true_df = df.iloc[pred_start : pred_end].copy()\n y_timestamp = y_true_df['date']\n \n print(f\"Backtesting: {x_df['date'].iloc[0].date()} to {y_timestamp.iloc[-1].date()}\")\n \n pred_df = predictor.predict(\n df=x_df[['open', 'high', 'low', 'close', 'volume']],\n x_timestamp=x_df['date'],\n y_timestamp=y_timestamp,\n pred_len=actual_pred_len,\n T=1.0, top_p=0.9, sample_count=1\n )\n \n render_comparison_chart(x_df, y_true_df, pred_df, f\"Backtest: {TICKER} K-Line Comparison\")\n\ndef run_forecast(df, predictor, lookback, pred_len):\n if len(df) \u003c lookback: return\n x_df = df.iloc[-lookback:].copy()\n last_date = x_df['date'].iloc[-1]\n future_dates = pd.date_range(start=last_date + BusinessDay(1), periods=pred_len, freq='B')\n future_dates = pd.Series(future_dates)\n \n print(f\"Forecasting: Starting from {future_dates.iloc[0].date()}\")\n \n pred_df = predictor.predict(\n df=x_df[['open', 'high', 'low', 'close', 'volume']],\n x_timestamp=x_df['date'],\n y_timestamp=future_dates,\n pred_len=pred_len,\n T=1.0, top_p=0.9, sample_count=1\n )\n \n render_comparison_chart(x_df, None, pred_df, f\"Forecast: {TICKER} Future K-Line\")\n\nif __name__ == \"__main__\":\n LOOKBACK = 20\n PRED_LEN = 10\n TICKER = '002111'\n \n pred_model = load_predictor()\n stock_data = load_data(TICKER)\n \n total_rows = len(stock_data)\n backtest_start = max(0, total_rows - LOOKBACK - PRED_LEN - 10) # Leave some space to see trend\n \n print(\"\\n--- Running Backtest ---\")\n run_backtest(stock_data, pred_model, LOOKBACK, PRED_LEN, start_index=backtest_start)\n \n print(\"\\n--- Running Forecast ---\")\n run_forecast(stock_data, pred_model, LOOKBACK, PRED_LEN)","content_type":"text/x-python; charset=utf-8","language":"python","size":7821,"content_sha256":"83139d1249d2d797969be1207d0dab9f59f9d2563f28799da83b3d0918b39952"},{"filename":"scripts/utils/predictor/model/__init__.py","content":"from .kronos import KronosTokenizer, Kronos, KronosPredictor\n\nmodel_dict = {\n 'kronos_tokenizer': KronosTokenizer,\n 'kronos': Kronos,\n 'kronos_predictor': KronosPredictor\n}\n\n\ndef get_model_class(model_name):\n if model_name in model_dict:\n return model_dict[model_name]\n else:\n print(f\"Model {model_name} not found in model_dict\")\n raise NotImplementedError\n\n","content_type":"text/x-python; charset=utf-8","language":"python","size":394,"content_sha256":"14193881034a621332145ee900e18cda79f72f65b2a2dc98059f8da5305ed336"},{"filename":"scripts/utils/predictor/model/kronos.py","content":"import numpy as np\nimport pandas as pd\nimport torch\nfrom huggingface_hub import PyTorchModelHubMixin\nimport sys\n\nfrom tqdm import trange\n\nsys.path.append(\"../\")\nfrom model.module import *\n\n\nclass KronosTokenizer(nn.Module, PyTorchModelHubMixin):\n \"\"\"\n KronosTokenizer module for tokenizing input data using a hybrid quantization approach.\n\n This tokenizer utilizes a combination of encoder and decoder Transformer blocks\n along with the Binary Spherical Quantization (BSQuantizer) to compress and decompress input data.\n\n Args:\n d_in (int): Input dimension.\n d_model (int): Model dimension.\n n_heads (int): Number of attention heads.\n ff_dim (int): Feed-forward dimension.\n n_enc_layers (int): Number of encoder layers.\n n_dec_layers (int): Number of decoder layers.\n ffn_dropout_p (float): Dropout probability for feed-forward networks.\n attn_dropout_p (float): Dropout probability for attention mechanisms.\n resid_dropout_p (float): Dropout probability for residual connections.\n s1_bits (int): Number of bits for the pre token in BSQuantizer.\n s2_bits (int): Number of bits for the post token in BSQuantizer.\n beta (float): Beta parameter for BSQuantizer.\n gamma0 (float): Gamma0 parameter for BSQuantizer.\n gamma (float): Gamma parameter for BSQuantizer.\n zeta (float): Zeta parameter for BSQuantizer.\n group_size (int): Group size parameter for BSQuantizer.\n\n \"\"\"\n\n def __init__(self, d_in, d_model, n_heads, ff_dim, n_enc_layers, n_dec_layers, ffn_dropout_p, attn_dropout_p, resid_dropout_p, s1_bits, s2_bits, beta, gamma0, gamma, zeta, group_size):\n\n super().__init__()\n self.d_in = d_in\n self.d_model = d_model\n self.n_heads = n_heads\n self.ff_dim = ff_dim\n self.enc_layers = n_enc_layers\n self.dec_layers = n_dec_layers\n self.ffn_dropout_p = ffn_dropout_p\n self.attn_dropout_p = attn_dropout_p\n self.resid_dropout_p = resid_dropout_p\n\n self.s1_bits = s1_bits\n self.s2_bits = s2_bits\n self.codebook_dim = s1_bits + s2_bits # Total dimension of the codebook after quantization\n self.embed = nn.Linear(self.d_in, self.d_model)\n self.head = nn.Linear(self.d_model, self.d_in)\n\n # Encoder Transformer Blocks\n self.encoder = nn.ModuleList([\n TransformerBlock(self.d_model, self.n_heads, self.ff_dim, self.ffn_dropout_p, self.attn_dropout_p, self.resid_dropout_p)\n for _ in range(self.enc_layers - 1)\n ])\n # Decoder Transformer Blocks\n self.decoder = nn.ModuleList([\n TransformerBlock(self.d_model, self.n_heads, self.ff_dim, self.ffn_dropout_p, self.attn_dropout_p, self.resid_dropout_p)\n for _ in range(self.dec_layers - 1)\n ])\n self.quant_embed = nn.Linear(in_features=self.d_model, out_features=self.codebook_dim) # Linear layer before quantization\n self.post_quant_embed_pre = nn.Linear(in_features=self.s1_bits, out_features=self.d_model) # Linear layer after quantization (pre part - s1 bits)\n self.post_quant_embed = nn.Linear(in_features=self.codebook_dim, out_features=self.d_model) # Linear layer after quantization (full codebook)\n self.tokenizer = BSQuantizer(self.s1_bits, self.s2_bits, beta, gamma0, gamma, zeta, group_size) # BSQuantizer module\n\n def forward(self, x):\n \"\"\"\n Forward pass of the KronosTokenizer.\n\n Args:\n x (torch.Tensor): Input tensor of shape (batch_size, seq_len, d_in).\n\n Returns:\n tuple: A tuple containing:\n - tuple: (z_pre, z) - Reconstructed outputs from decoder with s1_bits and full codebook respectively,\n both of shape (batch_size, seq_len, d_in).\n - torch.Tensor: bsq_loss - Loss from the BSQuantizer.\n - torch.Tensor: quantized - Quantized representation from BSQuantizer.\n - torch.Tensor: z_indices - Indices from the BSQuantizer.\n \"\"\"\n z = self.embed(x)\n\n for layer in self.encoder:\n z = layer(z)\n\n z = self.quant_embed(z) # (B, T, codebook)\n\n bsq_loss, quantized, z_indices = self.tokenizer(z)\n\n quantized_pre = quantized[:, :, :self.s1_bits] # Extract the first part of quantized representation (s1_bits)\n z_pre = self.post_quant_embed_pre(quantized_pre)\n\n z = self.post_quant_embed(quantized)\n\n # Decoder layers (for pre part - s1 bits)\n for layer in self.decoder:\n z_pre = layer(z_pre)\n z_pre = self.head(z_pre)\n\n # Decoder layers (for full codebook)\n for layer in self.decoder:\n z = layer(z)\n z = self.head(z)\n\n return (z_pre, z), bsq_loss, quantized, z_indices\n\n def indices_to_bits(self, x, half=False):\n \"\"\"\n Converts indices to bit representations and scales them.\n\n Args:\n x (torch.Tensor): Indices tensor.\n half (bool, optional): Whether to process only half of the codebook dimension. Defaults to False.\n\n Returns:\n torch.Tensor: Bit representation tensor.\n \"\"\"\n if half:\n x1 = x[0] # Assuming x is a tuple of indices if half is True\n x2 = x[1]\n mask = 2 ** torch.arange(self.codebook_dim//2, device=x1.device, dtype=torch.long) # Create a mask for bit extraction\n x1 = (x1.unsqueeze(-1) & mask) != 0 # Extract bits for the first half\n x2 = (x2.unsqueeze(-1) & mask) != 0 # Extract bits for the second half\n x = torch.cat([x1, x2], dim=-1) # Concatenate the bit representations\n else:\n mask = 2 ** torch.arange(self.codebook_dim, device=x.device, dtype=torch.long) # Create a mask for bit extraction\n x = (x.unsqueeze(-1) & mask) != 0 # Extract bits\n\n x = x.float() * 2 - 1 # Convert boolean to bipolar (-1, 1)\n q_scale = 1. / (self.codebook_dim ** 0.5) # Scaling factor\n x = x * q_scale\n return x\n\n def encode(self, x, half=False):\n \"\"\"\n Encodes the input data into quantized indices.\n\n Args:\n x (torch.Tensor): Input tensor of shape (batch_size, seq_len, d_in).\n half (bool, optional): Whether to use half quantization in BSQuantizer. Defaults to False.\n\n Returns:\n torch.Tensor: Quantized indices from BSQuantizer.\n \"\"\"\n z = self.embed(x)\n for layer in self.encoder:\n z = layer(z)\n z = self.quant_embed(z)\n\n bsq_loss, quantized, z_indices = self.tokenizer(z, half=half, collect_metrics=False)\n return z_indices\n\n def decode(self, x, half=False):\n \"\"\"\n Decodes quantized indices back to the input data space.\n\n Args:\n x (torch.Tensor): Quantized indices tensor.\n half (bool, optional): Whether the indices were generated with half quantization. Defaults to False.\n\n Returns:\n torch.Tensor: Reconstructed output tensor of shape (batch_size, seq_len, d_in).\n \"\"\"\n quantized = self.indices_to_bits(x, half)\n z = self.post_quant_embed(quantized)\n for layer in self.decoder:\n z = layer(z)\n z = self.head(z)\n return z\n\n\nclass Kronos(nn.Module, PyTorchModelHubMixin):\n \"\"\"\n Kronos Model.\n\n Args:\n s1_bits (int): Number of bits for pre tokens.\n s2_bits (int): Number of bits for post tokens.\n n_layers (int): Number of Transformer blocks.\n d_model (int): Dimension of the model's embeddings and hidden states.\n n_heads (int): Number of attention heads in the MultiheadAttention layers.\n ff_dim (int): Dimension of the feedforward network in the Transformer blocks.\n ffn_dropout_p (float): Dropout probability for the feedforward network.\n attn_dropout_p (float): Dropout probability for the attention layers.\n resid_dropout_p (float): Dropout probability for residual connections.\n token_dropout_p (float): Dropout probability for token embeddings.\n learn_te (bool): Whether to use learnable temporal embeddings.\n \"\"\"\n\n def __init__(self, s1_bits, s2_bits, n_layers, d_model, n_heads, ff_dim, ffn_dropout_p, attn_dropout_p, resid_dropout_p, token_dropout_p, learn_te, news_dim=None):\n super().__init__()\n self.s1_bits = s1_bits\n self.s2_bits = s2_bits\n self.n_layers = n_layers\n self.d_model = d_model\n self.n_heads = n_heads\n self.learn_te = learn_te\n self.ff_dim = ff_dim\n self.ffn_dropout_p = ffn_dropout_p\n self.attn_dropout_p = attn_dropout_p\n self.resid_dropout_p = resid_dropout_p\n self.token_dropout_p = token_dropout_p\n self.news_dim = news_dim\n\n self.s1_vocab_size = 2 ** self.s1_bits\n self.token_drop = nn.Dropout(self.token_dropout_p)\n self.embedding = HierarchicalEmbedding(self.s1_bits, self.s2_bits, self.d_model)\n self.time_emb = TemporalEmbedding(self.d_model, self.learn_te)\n self.transformer = nn.ModuleList([\n TransformerBlock(self.d_model, self.n_heads, self.ff_dim, self.ffn_dropout_p, self.attn_dropout_p, self.resid_dropout_p)\n for _ in range(self.n_layers)\n ])\n self.norm = RMSNorm(self.d_model)\n self.dep_layer = DependencyAwareLayer(self.d_model)\n self.head = DualHead(self.s1_bits, self.s2_bits, self.d_model)\n\n if self.news_dim is not None:\n self.news_proj = nn.Linear(self.news_dim, self.d_model)\n else:\n self.news_proj = None\n\n self.apply(self._init_weights)\n\n def _init_weights(self, module):\n\n if isinstance(module, nn.Linear):\n nn.init.xavier_normal_(module.weight)\n if module.bias is not None:\n nn.init.zeros_(module.bias)\n elif isinstance(module, nn.Embedding):\n nn.init.normal_(module.weight, mean=0, std=self.embedding.d_model ** -0.5)\n elif isinstance(module, nn.LayerNorm):\n nn.init.ones_(module.weight)\n nn.init.zeros_(module.bias)\n elif isinstance(module, RMSNorm):\n nn.init.ones_(module.weight)\n\n def forward(self, s1_ids, s2_ids, stamp=None, padding_mask=None, use_teacher_forcing=False, s1_targets=None, news_emb=None):\n \"\"\"\n Args:\n s1_ids (torch.Tensor): Input tensor of s1 token IDs. Shape: [batch_size, seq_len]\n s2_ids (torch.Tensor): Input tensor of s2 token IDs. Shape: [batch_size, seq_len]\n stamp (torch.Tensor, optional): Temporal stamp tensor. Shape: [batch_size, seq_len]. Defaults to None.\n padding_mask (torch.Tensor, optional): Mask for padding tokens. Shape: [batch_size, seq_len]. Defaults to None.\n use_teacher_forcing (bool, optional): Whether to use teacher forcing for s1 decoding. Defaults to False.\n s1_targets (torch.Tensor, optional): Target s1 token IDs for teacher forcing. Shape: [batch_size, seq_len]. Defaults to None.\n news_emb (torch.Tensor, optional): News embedding tensor. Shape: [batch_size, news_dim]. Defaults to None.\n\n Returns:\n Tuple[torch.Tensor, torch.Tensor]:\n - s1 logits: Logits for s1 token predictions. Shape: [batch_size, seq_len, s1_vocab_size]\n - s2_logits: Logits for s2 token predictions, conditioned on s1. Shape: [batch_size, seq_len, s2_vocab_size]\n \"\"\"\n x = self.embedding([s1_ids, s2_ids])\n if stamp is not None:\n time_embedding = self.time_emb(stamp)\n x = x + time_embedding\n x = self.token_drop(x)\n\n for layer in self.transformer:\n x = layer(x, key_padding_mask=padding_mask)\n\n x = self.norm(x)\n\n if news_emb is not None and self.news_proj is not None:\n news_bias = self.news_proj(news_emb).unsqueeze(1) # [B, 1, d_model]\n x = x + news_bias\n\n s1_logits = self.head(x)\n\n if use_teacher_forcing:\n sibling_embed = self.embedding.emb_s1(s1_targets)\n else:\n s1_probs = F.softmax(s1_logits.detach(), dim=-1)\n sample_s1_ids = torch.multinomial(s1_probs.view(-1, self.s1_vocab_size), 1).view(s1_ids.shape)\n sibling_embed = self.embedding.emb_s1(sample_s1_ids)\n\n x2 = self.dep_layer(x, sibling_embed, key_padding_mask=padding_mask) # Dependency Aware Layer: Condition on s1 embeddings\n s2_logits = self.head.cond_forward(x2)\n return s1_logits, s2_logits\n\n def decode_s1(self, s1_ids, s2_ids, stamp=None, padding_mask=None, news_emb=None):\n \"\"\"\n Decodes only the s1 tokens.\n\n This method performs a forward pass to predict only s1 tokens. It returns the s1 logits\n and the context representation from the Transformer, which can be used for subsequent s2 decoding.\n\n Args:\n s1_ids (torch.Tensor): Input tensor of s1 token IDs. Shape: [batch_size, seq_len]\n s2_ids (torch.Tensor): Input tensor of s2 token IDs. Shape: [batch_size, seq_len]\n stamp (torch.Tensor, optional): Temporal stamp tensor. Shape: [batch_size, seq_len]. Defaults to None.\n padding_mask (torch.Tensor, optional): Mask for padding tokens. Shape: [batch_size, seq_len]. Defaults to None.\n news_emb (torch.Tensor, optional): News embedding tensor. Shape: [batch_size, news_dim]. Defaults to None.\n\n Returns:\n Tuple[torch.Tensor, torch.Tensor]:\n - s1 logits: Logits for s1 token predictions. Shape: [batch_size, seq_len, s1_vocab_size]\n - context: Context representation from the Transformer. Shape: [batch_size, seq_len, d_model]\n \"\"\"\n x = self.embedding([s1_ids, s2_ids])\n if stamp is not None:\n time_embedding = self.time_emb(stamp)\n x = x + time_embedding\n x = self.token_drop(x)\n\n for layer in self.transformer:\n x = layer(x, key_padding_mask=padding_mask)\n\n x = self.norm(x)\n\n if news_emb is not None and self.news_proj is not None:\n news_bias = self.news_proj(news_emb).unsqueeze(1) # [B, 1, d_model]\n x = x + news_bias\n\n s1_logits = self.head(x)\n return s1_logits, x\n\n def decode_s2(self, context, s1_ids, padding_mask=None):\n \"\"\"\n Decodes the s2 tokens, conditioned on the context and s1 tokens.\n\n This method decodes s2 tokens based on a pre-computed context representation (typically from `decode_s1`)\n and the s1 token IDs. It uses the dependency-aware layer and the conditional s2 head to predict s2 tokens.\n\n Args:\n context (torch.Tensor): Context representation from the transformer (output of decode_s1).\n Shape: [batch_size, seq_len, d_model]\n s1_ids (torch.torch.Tensor): Input tensor of s1 token IDs. Shape: [batch_size, seq_len]\n padding_mask (torch.Tensor, optional): Mask for padding tokens. Shape: [batch_size, seq_len]. Defaults to None.\n\n Returns:\n torch.Tensor: s2 logits. Shape: [batch_size, seq_len, s2_vocab_size]\n \"\"\"\n sibling_embed = self.embedding.emb_s1(s1_ids)\n x2 = self.dep_layer(context, sibling_embed, key_padding_mask=padding_mask)\n return self.head.cond_forward(x2)\n\n\ndef top_k_top_p_filtering(\n logits,\n top_k: int = 0,\n top_p: float = 1.0,\n filter_value: float = -float(\"Inf\"),\n min_tokens_to_keep: int = 1,\n):\n \"\"\"Filter a distribution of logits using top-k and/or nucleus (top-p) filtering\n Args:\n logits: logits distribution shape (batch size, vocabulary size)\n if top_k > 0: keep only top k tokens with highest probability (top-k filtering).\n if top_p \u003c 1.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).\n Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)\n Make sure we keep at least min_tokens_to_keep per batch example in the output\n From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317\n \"\"\"\n if top_k > 0:\n top_k = min(max(top_k, min_tokens_to_keep), logits.size(-1)) # Safety check\n # Remove all tokens with a probability less than the last token of the top-k\n indices_to_remove = logits \u003c torch.topk(logits, top_k)[0][..., -1, None]\n logits[indices_to_remove] = filter_value\n return logits\n\n if top_p \u003c 1.0:\n sorted_logits, sorted_indices = torch.sort(logits, descending=True)\n cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)\n\n # Remove tokens with cumulative probability above the threshold (token with 0 are kept)\n sorted_indices_to_remove = cumulative_probs > top_p\n if min_tokens_to_keep > 1:\n # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)\n sorted_indices_to_remove[..., :min_tokens_to_keep] = 0\n # Shift the indices to the right to keep also the first token above the threshold\n sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()\n sorted_indices_to_remove[..., 0] = 0\n\n # scatter sorted tensors to original indexing\n indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)\n logits[indices_to_remove] = filter_value\n return logits\n\n\ndef sample_from_logits(logits, temperature=1.0, top_k=None, top_p=None, sample_logits=True):\n logits = logits / temperature\n if top_k is not None or top_p is not None:\n if top_k > 0 or top_p \u003c 1.0:\n logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)\n\n probs = F.softmax(logits, dim=-1)\n\n if not sample_logits:\n _, x = top_k(probs, k=1, dim=-1)\n else:\n x = torch.multinomial(probs, num_samples=1)\n\n return x\n\n\ndef auto_regressive_inference(tokenizer, model, x, x_stamp, y_stamp, max_context, pred_len, clip=5, T=1.0, top_k=0, top_p=0.99, sample_count=5, verbose=False, news_emb=None):\n with torch.no_grad():\n x = torch.clip(x, -clip, clip)\n\n device = x.device\n x = x.unsqueeze(1).repeat(1, sample_count, 1, 1).reshape(-1, x.size(1), x.size(2)).to(device)\n x_stamp = x_stamp.unsqueeze(1).repeat(1, sample_count, 1, 1).reshape(-1, x_stamp.size(1), x_stamp.size(2)).to(device)\n y_stamp = y_stamp.unsqueeze(1).repeat(1, sample_count, 1, 1).reshape(-1, y_stamp.size(1), y_stamp.size(2)).to(device)\n\n x_token = tokenizer.encode(x, half=True)\n \n initial_seq_len = x.size(1)\n batch_size = x_token[0].size(0)\n total_seq_len = initial_seq_len + pred_len\n full_stamp = torch.cat([x_stamp, y_stamp], dim=1)\n\n generated_pre = x_token[0].new_empty(batch_size, pred_len)\n generated_post = x_token[1].new_empty(batch_size, pred_len)\n\n pre_buffer = x_token[0].new_zeros(batch_size, max_context)\n post_buffer = x_token[1].new_zeros(batch_size, max_context)\n buffer_len = min(initial_seq_len, max_context)\n if buffer_len > 0:\n start_idx = max(0, initial_seq_len - max_context)\n pre_buffer[:, :buffer_len] = x_token[0][:, start_idx:start_idx + buffer_len]\n post_buffer[:, :buffer_len] = x_token[1][:, start_idx:start_idx + buffer_len]\n\n if verbose:\n ran = trange\n else:\n ran = range\n for i in ran(pred_len):\n current_seq_len = initial_seq_len + i\n window_len = min(current_seq_len, max_context)\n\n if current_seq_len \u003c= max_context:\n input_tokens = [\n pre_buffer[:, :window_len],\n post_buffer[:, :window_len]\n ]\n else:\n input_tokens = [pre_buffer, post_buffer]\n\n context_end = current_seq_len\n context_start = max(0, context_end - max_context)\n current_stamp = full_stamp[:, context_start:context_end, :].contiguous()\n\n s1_logits, context = model.decode_s1(input_tokens[0], input_tokens[1], current_stamp, news_emb=news_emb)\n s1_logits = s1_logits[:, -1, :]\n sample_pre = sample_from_logits(s1_logits, temperature=T, top_k=top_k, top_p=top_p, sample_logits=True)\n\n s2_logits = model.decode_s2(context, sample_pre)\n s2_logits = s2_logits[:, -1, :]\n sample_post = sample_from_logits(s2_logits, temperature=T, top_k=top_k, top_p=top_p, sample_logits=True)\n\n generated_pre[:, i] = sample_pre.squeeze(-1)\n generated_post[:, i] = sample_post.squeeze(-1)\n\n if current_seq_len \u003c max_context:\n pre_buffer[:, current_seq_len] = sample_pre.squeeze(-1)\n post_buffer[:, current_seq_len] = sample_post.squeeze(-1)\n else:\n pre_buffer.copy_(torch.roll(pre_buffer, shifts=-1, dims=1))\n post_buffer.copy_(torch.roll(post_buffer, shifts=-1, dims=1))\n pre_buffer[:, -1] = sample_pre.squeeze(-1)\n post_buffer[:, -1] = sample_post.squeeze(-1)\n\n full_pre = torch.cat([x_token[0], generated_pre], dim=1)\n full_post = torch.cat([x_token[1], generated_post], dim=1)\n\n context_start = max(0, total_seq_len - max_context)\n input_tokens = [\n full_pre[:, context_start:total_seq_len].contiguous(),\n full_post[:, context_start:total_seq_len].contiguous()\n ]\n z = tokenizer.decode(input_tokens, half=True)\n z = z.reshape(-1, sample_count, z.size(1), z.size(2))\n preds = z.cpu().numpy()\n preds = np.mean(preds, axis=1)\n\n return preds\n\n\ndef calc_time_stamps(x_timestamp):\n time_df = pd.DataFrame()\n time_df['minute'] = x_timestamp.dt.minute\n time_df['hour'] = x_timestamp.dt.hour\n time_df['weekday'] = x_timestamp.dt.weekday\n time_df['day'] = x_timestamp.dt.day\n time_df['month'] = x_timestamp.dt.month\n return time_df\n\n\nclass KronosPredictor:\n\n def __init__(self, model, tokenizer, device=\"cuda:0\", max_context=512, clip=5):\n self.tokenizer = tokenizer\n self.model = model\n self.max_context = max_context\n self.clip = clip\n self.price_cols = ['open', 'high', 'low', 'close']\n self.vol_col = 'volume'\n self.amt_vol = 'amount'\n self.time_cols = ['minute', 'hour', 'weekday', 'day', 'month']\n self.device = device\n\n self.tokenizer = self.tokenizer.to(self.device)\n self.model = self.model.to(self.device)\n\n def generate(self, x, x_stamp, y_stamp, pred_len, T, top_k, top_p, sample_count, verbose, news_emb=None):\n\n x_tensor = torch.from_numpy(np.array(x).astype(np.float32)).to(self.device)\n x_stamp_tensor = torch.from_numpy(np.array(x_stamp).astype(np.float32)).to(self.device)\n y_stamp_tensor = torch.from_numpy(np.array(y_stamp).astype(np.float32)).to(self.device)\n\n preds = auto_regressive_inference(self.tokenizer, self.model, x_tensor, x_stamp_tensor, y_stamp_tensor, self.max_context, pred_len,\n self.clip, T, top_k, top_p, sample_count, verbose, news_emb=news_emb)\n preds = preds[:, -pred_len:, :]\n return preds\n\n def predict(self, df, x_timestamp, y_timestamp, pred_len, T=1.0, top_k=0, top_p=0.9, sample_count=1, verbose=True, news_emb=None):\n\n if not isinstance(df, pd.DataFrame):\n raise ValueError(\"Input must be a pandas DataFrame.\")\n\n if not all(col in df.columns for col in self.price_cols):\n raise ValueError(f\"Price columns {self.price_cols} not found in DataFrame.\")\n\n df = df.copy()\n if self.vol_col not in df.columns:\n df[self.vol_col] = 0.0 # Fill missing volume with zeros\n df[self.amt_vol] = 0.0 # Fill missing amount with zeros\n if self.amt_vol not in df.columns and self.vol_col in df.columns:\n df[self.amt_vol] = df[self.vol_col] * df[self.price_cols].mean(axis=1)\n\n if df[self.price_cols + [self.vol_col, self.amt_vol]].isnull().values.any():\n raise ValueError(\"Input DataFrame contains NaN values in price or volume columns.\")\n\n x_time_df = calc_time_stamps(x_timestamp)\n y_time_df = calc_time_stamps(y_timestamp)\n\n x = df[self.price_cols + [self.vol_col, self.amt_vol]].values.astype(np.float32)\n x_stamp = x_time_df.values.astype(np.float32)\n y_stamp = y_time_df.values.astype(np.float32)\n\n x_mean, x_std = np.mean(x, axis=0), np.std(x, axis=0)\n\n x = (x - x_mean) / (x_std + 1e-5)\n x = np.clip(x, -self.clip, self.clip)\n\n x = x[np.newaxis, :]\n x_stamp = x_stamp[np.newaxis, :]\n y_stamp = y_stamp[np.newaxis, :]\n\n if news_emb is not None:\n news_emb_tensor = torch.from_numpy(np.array(news_emb).astype(np.float32)).to(self.device)\n # Ensure batch dimension for news_emb if only one sample\n if news_emb_tensor.ndim == 1:\n news_emb_tensor = news_emb_tensor.unsqueeze(0)\n else:\n news_emb_tensor = None\n\n preds = self.generate(x, x_stamp, y_stamp, pred_len, T, top_k, top_p, sample_count, verbose, news_emb=news_emb_tensor)\n\n preds = preds.squeeze(0)\n preds = preds * (x_std + 1e-5) + x_mean\n\n pred_df = pd.DataFrame(preds, columns=self.price_cols + [self.vol_col, self.amt_vol], index=y_timestamp)\n return pred_df\n\n\n def predict_batch(self, df_list, x_timestamp_list, y_timestamp_list, pred_len, T=1.0, top_k=0, top_p=0.9, sample_count=1, verbose=True):\n \"\"\"\n Perform parallel (batch) prediction on multiple time series. All series must have the same historical length and prediction length (pred_len).\n\n Args:\n df_list (List[pd.DataFrame]): List of input DataFrames, each containing price columns and optional volume/amount columns.\n x_timestamp_list (List[pd.DatetimeIndex or Series]): List of timestamps corresponding to historical data, length should match the number of rows in each DataFrame.\n y_timestamp_list (List[pd.DatetimeIndex or Series]): List of future prediction timestamps, length should equal pred_len.\n pred_len (int): Number of prediction steps.\n T (float): Sampling temperature.\n top_k (int): Top-k filtering threshold.\n top_p (float): Top-p (nucleus sampling) threshold.\n sample_count (int): Number of parallel samples per series, automatically averaged internally.\n verbose (bool): Whether to display autoregressive progress.\n\n Returns:\n List[pd.DataFrame]: List of prediction results in the same order as input, each DataFrame contains\n `open, high, low, close, volume, amount` columns, indexed by corresponding `y_timestamp`.\n \"\"\"\n # Basic validation\n if not isinstance(df_list, (list, tuple)) or not isinstance(x_timestamp_list, (list, tuple)) or not isinstance(y_timestamp_list, (list, tuple)):\n raise ValueError(\"df_list, x_timestamp_list, y_timestamp_list must be list or tuple types.\")\n if not (len(df_list) == len(x_timestamp_list) == len(y_timestamp_list)):\n raise ValueError(\"df_list, x_timestamp_list, y_timestamp_list must have consistent lengths.\")\n\n num_series = len(df_list)\n\n x_list = []\n x_stamp_list = []\n y_stamp_list = []\n means = []\n stds = []\n seq_lens = []\n y_lens = []\n\n for i in range(num_series):\n df = df_list[i]\n if not isinstance(df, pd.DataFrame):\n raise ValueError(f\"Input at index {i} is not a pandas DataFrame.\")\n if not all(col in df.columns for col in self.price_cols):\n raise ValueError(f\"DataFrame at index {i} is missing price columns {self.price_cols}.\")\n\n df = df.copy()\n if self.vol_col not in df.columns:\n df[self.vol_col] = 0.0\n df[self.amt_vol] = 0.0\n if self.amt_vol not in df.columns and self.vol_col in df.columns:\n df[self.amt_vol] = df[self.vol_col] * df[self.price_cols].mean(axis=1)\n\n if df[self.price_cols + [self.vol_col, self.amt_vol]].isnull().values.any():\n raise ValueError(f\"DataFrame at index {i} contains NaN values in price or volume columns.\")\n\n x_timestamp = x_timestamp_list[i]\n y_timestamp = y_timestamp_list[i]\n\n x_time_df = calc_time_stamps(x_timestamp)\n y_time_df = calc_time_stamps(y_timestamp)\n\n x = df[self.price_cols + [self.vol_col, self.amt_vol]].values.astype(np.float32)\n x_stamp = x_time_df.values.astype(np.float32)\n y_stamp = y_time_df.values.astype(np.float32)\n\n if x.shape[0] != x_stamp.shape[0]:\n raise ValueError(f\"Inconsistent lengths at index {i}: x has {x.shape[0]} vs x_stamp has {x_stamp.shape[0]}.\")\n if y_stamp.shape[0] != pred_len:\n raise ValueError(f\"y_timestamp length at index {i} should equal pred_len={pred_len}, got {y_stamp.shape[0]}.\")\n\n x_mean, x_std = np.mean(x, axis=0), np.std(x, axis=0)\n x_norm = (x - x_mean) / (x_std + 1e-5)\n x_norm = np.clip(x_norm, -self.clip, self.clip)\n\n x_list.append(x_norm)\n x_stamp_list.append(x_stamp)\n y_stamp_list.append(y_stamp)\n means.append(x_mean)\n stds.append(x_std)\n\n seq_lens.append(x_norm.shape[0])\n y_lens.append(y_stamp.shape[0])\n\n # Require all series to have consistent historical and prediction lengths for batch processing\n if len(set(seq_lens)) != 1:\n raise ValueError(f\"Parallel prediction requires all series to have consistent historical lengths, got: {seq_lens}\")\n if len(set(y_lens)) != 1:\n raise ValueError(f\"Parallel prediction requires all series to have consistent prediction lengths, got: {y_lens}\")\n\n x_batch = np.stack(x_list, axis=0).astype(np.float32) # (B, seq_len, feat)\n x_stamp_batch = np.stack(x_stamp_list, axis=0).astype(np.float32) # (B, seq_len, time_feat)\n y_stamp_batch = np.stack(y_stamp_list, axis=0).astype(np.float32) # (B, pred_len, time_feat)\n\n preds = self.generate(x_batch, x_stamp_batch, y_stamp_batch, pred_len, T, top_k, top_p, sample_count, verbose)\n # preds: (B, pred_len, feat)\n\n pred_dfs = []\n for i in range(num_series):\n preds_i = preds[i] * (stds[i] + 1e-5) + means[i]\n pred_df = pd.DataFrame(preds_i, columns=self.price_cols + [self.vol_col, self.amt_vol], index=y_timestamp_list[i])\n pred_dfs.append(pred_df)\n\n return pred_dfs\n","content_type":"text/x-python; charset=utf-8","language":"python","size":31093,"content_sha256":"b339c8685f40322663eb6c3f0d56dc91e64b65fae1f9cc22e414d6c76b9c7b78"},{"filename":"scripts/utils/predictor/model/module.py","content":"import math\n\nfrom einops import rearrange, reduce\nimport torch\nimport torch.nn as nn\nfrom torch.autograd import Function\nimport torch.nn.functional as F\n\n\nclass DifferentiableEntropyFunction(Function):\n @staticmethod\n def forward(ctx, zq, basis, K, eps):\n zb = (zq + 1) / 2\n zi = ((zb * basis).sum(-1)).to(torch.int64)\n cnt = torch.scatter_reduce(torch.zeros(2 ** K, device=zq.device, dtype=zq.dtype),\n 0,\n zi.flatten(),\n torch.ones_like(zi.flatten()).to(zq.dtype),\n 'sum')\n prob = (cnt + eps) / (cnt + eps).sum()\n H = -(prob * torch.log(prob)).sum()\n ctx.save_for_backward(zq, zi, prob)\n ctx.K = K\n return H\n\n @staticmethod\n def backward(ctx, grad_output):\n zq, zi, prob = ctx.saved_tensors\n grad_array = -grad_output * (torch.log(prob) + 1) / zi.numel() / ctx.K\n reord_grad = grad_array[zi.flatten()].reshape(zi.shape)\n grad_input = reord_grad.unsqueeze(-1) * zq\n return grad_input, None, None, None, None\n\n\ndef codebook_entropy(zq, basis, K, eps=1e-4):\n return DifferentiableEntropyFunction.apply(zq, basis, K, eps)\n\n\nclass BinarySphericalQuantizer(nn.Module):\n def __init__(self, embed_dim, beta, gamma0, gamma, zeta,\n input_format='bchw',\n soft_entropy=True, group_size=9,\n persample_entropy_compute='analytical',\n cb_entropy_compute='group',\n l2_norm=True,\n inv_temperature=1):\n \"\"\"\n Paper link: https://arxiv.org/pdf/2406.07548.pdf\n Here we use the official implementation of the BinarySphericalQuantizer.\n \"\"\"\n super().__init__()\n self.embed_dim = embed_dim\n self.beta = beta # loss weight for commit loss\n self.gamma0 = gamma0 # loss weight for entropy penalty\n self.gamma = gamma # loss weight for entropy penalty\n self.zeta = zeta # loss weight for entire entropy penalty\n self.input_format = input_format\n assert self.embed_dim % group_size == 0, \"embed_dim must be divisible by group_size\"\n self.num_groups = self.embed_dim // group_size\n self.group_size = group_size\n assert persample_entropy_compute in ['group', 'analytical'], \"persample_entropy_compute must be either 'group' or 'analytical'\"\n assert cb_entropy_compute in ['group', 'nce'], \"cb_entropy_compute must be either 'group' or 'nce'\"\n self.persample_entropy_compute = persample_entropy_compute\n self.cb_entropy_compute = cb_entropy_compute\n self.l2_norm = l2_norm\n self.inv_temperature = inv_temperature\n\n self.register_buffer('basis', 2 ** torch.arange(embed_dim - 1, -1, -1))\n self.register_buffer('group_basis', 2 ** torch.arange(group_size - 1, -1, -1))\n\n self.num_dimensions = 2 ** embed_dim\n self.bits_per_index = embed_dim\n\n # we only need to keep the codebook portion up to the group size\n # because we approximate the H loss with this subcode\n group_codes = torch.arange(2 ** self.group_size)\n group_codebook = self.indexes_to_codes(group_codes).float()[:, -group_size:]\n self.register_buffer('group_codebook', group_codebook, persistent=False)\n\n self.soft_entropy = soft_entropy # soft_entropy: Sec 3.2 of https://arxiv.org/pdf/1911.05894.pdf\n\n def quantize(self, z):\n assert z.shape[-1] == self.embed_dim, f\"Expected {self.embed_dim} dimensions, got {z.shape[-1]}\"\n\n zhat = torch.where(z > 0,\n torch.tensor(1, dtype=z.dtype, device=z.device),\n torch.tensor(-1, dtype=z.dtype, device=z.device))\n return z + (zhat - z).detach()\n\n def forward(self, z, collect_metrics=True):\n # if self.input_format == 'bchw':\n # z = rearrange(z, 'b c h w -> b h w c')\n zq = self.quantize(z)\n\n q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.\n\n zq = zq * q_scale\n\n if not collect_metrics:\n return zq, zq.new_zeros(()), {}\n\n indices = self.codes_to_indexes(zq.detach())\n group_indices = self.codes_to_group_indexes(zq.detach())\n if not self.training:\n used_codes = torch.unique(indices, return_counts=False)\n else:\n used_codes = None\n\n if self.soft_entropy:\n persample_entropy, cb_entropy, avg_prob = self.soft_entropy_loss(z)\n entropy_penalty = self.gamma0 * persample_entropy - self.gamma * cb_entropy\n else:\n zb_by_sample = ((zq + 1) / 2).reshape(z.shape[0], -1, z.shape[-1]).to(torch.float32)\n persample_entropy = self.get_hard_per_sample_entropy(zb_by_sample)\n cb_entropy = codebook_entropy(zq, self.basis, self.embed_dim)\n entropy_penalty = self.gamma0 * persample_entropy - self.gamma * cb_entropy\n\n # commit loss\n commit_loss = self.beta * torch.mean(((zq.detach() - z) ** 2).sum(dim=-1))\n\n # if self.input_format == 'bchw':\n # zq = rearrange(zq, 'b h w c -> b c h w')\n\n return (\n zq,\n commit_loss + self.zeta * entropy_penalty / self.inv_temperature,\n {\"H\": cb_entropy, \"used_codes\": used_codes, \"indices\": indices, \"group_indices\": group_indices,\n \"avg_prob\": avg_prob}\n )\n\n def soft_entropy_loss(self, z):\n # if we divide the code in subgroups of size group_size, the codebook will be of size 2 ** group_size\n # the sub-code is the last group_size bits of the full code\n group_code_book = self.group_codebook / (self.embed_dim ** 0.5 if self.l2_norm else 1)\n divided_z = rearrange(z, '... (g c) -> ... g c', c=self.group_size)\n\n # we calculate the distance between the divided_z and the codebook for each subgroup\n distance = - 2 * torch.einsum('... g c, d c ->... g d', divided_z, group_code_book)\n prob = (-distance * self.inv_temperature).softmax(dim=-1)\n if self.persample_entropy_compute == 'analytical':\n if self.l2_norm:\n p = torch.sigmoid(-4 * z / (self.embed_dim ** 0.5) * self.inv_temperature)\n else:\n p = torch.sigmoid(-4 * z * self.inv_temperature)\n prob = torch.stack([p, 1 - p], dim=-1)\n per_sample_entropy = self.get_entropy(prob, dim=-1, normalize=False).sum(dim=-1).mean()\n else:\n per_sample_entropy = self.get_entropy(prob, dim=-1, normalize=False).sum(dim=-1).mean()\n\n # macro average of the probability of each subgroup\n avg_prob = reduce(prob, '... g d ->g d', 'mean')\n codebook_entropy = self.get_entropy(avg_prob, dim=-1, normalize=False)\n\n # the approximation of the entropy is the sum of the entropy of each subgroup\n return per_sample_entropy, codebook_entropy.sum(), avg_prob\n\n def get_hard_per_sample_entropy(self, zb_by_sample):\n probs_per_dim = zb_by_sample.sum(1) / zb_by_sample.shape[1]\n persample_entropy = - probs_per_dim * torch.log(probs_per_dim + 1e-8) - (1 - probs_per_dim) * torch.log(1 - probs_per_dim + 1e-8)\n persample_entropy = persample_entropy.sum(-1)\n return persample_entropy.mean()\n\n def codes_to_indexes(self, zhat):\n \"\"\"Converts a `code` to an index in the codebook.\n Args:\n zhat: A tensor of shape (B, ..., C) containing the codes. must be in {-1, 1}\n \"\"\"\n assert zhat.shape[-1] == self.embed_dim, f\"Expected {self.embed_dim} dimensions, got {zhat.shape[-1]}\"\n return ((zhat + 1) / 2 * self.basis).sum(axis=-1).to(torch.int64)\n\n def codes_to_group_indexes(self, zhat):\n \"\"\"Converts a `code` to a list of indexes (in groups) in the codebook.\n Args:\n zhat: A tensor of shape (B, ..., C) containing the codes. must be in {-1, 1}\n \"\"\"\n zhat_in_group = rearrange(zhat, 'b ... (g c) -> b ... g c', c=self.group_size)\n return ((zhat_in_group + 1) / 2 * self.group_basis).sum(axis=-1).to(torch.int64)\n\n def indexes_to_codes(self, indices):\n \"\"\"Inverse of `indexes_to_codes`.\"\"\"\n indices = indices.unsqueeze(-1)\n codes_non_centered = torch.remainder(\n torch.floor_divide(indices, self.basis), 2\n )\n return codes_non_centered * 2 - 1\n\n def group_indexes_to_codes(self, group_indices):\n \"\"\"Inverse of `group_indexes_to_codes`.\"\"\"\n group_indices = group_indices.unsqueeze(-1)\n codes_non_centered = torch.remainder(\n torch.floor_divide(group_indices, self.group_basis), 2\n )\n codes_non_centered = rearrange(codes_non_centered, 'b ... g c -> b ... (g c)')\n return codes_non_centered * 2 - 1\n\n def get_entropy(self, count, dim=-1, eps=1e-4, normalize=True):\n if normalize:\n probs = (count + eps) / (count + eps).sum(dim=dim, keepdim=True)\n else:\n probs = count\n H = -(probs * torch.log(probs + 1e-8)).sum(dim=dim)\n return H\n\n def get_group_codebook_entry(self, group_indices):\n z_q = self.group_indexes_to_codes(group_indices)\n q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.\n z_q = z_q * q_scale\n if self.input_format == 'bchw':\n h, w = int(z_q.shape[1] ** 0.5)\n assert h * w == z_q.shape[1], 'Invalid sequence length'\n z_q = rearrange(z_q, 'b (h w) c -> b c h w', h=h)\n return z_q\n\n def get_codebook_entry(self, indices):\n z_q = self.indexes_to_codes(indices)\n q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.\n z_q = z_q * q_scale\n if self.input_format == 'bchw':\n h, w = int(z_q.shape[1] ** 0.5)\n assert h * w == z_q.shape[1], 'Invalid sequence length'\n z_q = rearrange(z_q, 'b (h w) c -> b c h w', h=h)\n return z_q\n\n\nclass BSQuantizer(nn.Module):\n\n def __init__(self, s1_bits, s2_bits, beta, gamma0, gamma, zeta, group_size):\n super().__init__()\n self.codebook_dim = s1_bits + s2_bits\n self.s1_bits = s1_bits\n self.s2_bits = s2_bits\n self.bsq = BinarySphericalQuantizer(self.codebook_dim, beta, gamma0, gamma, zeta, group_size=group_size)\n\n def bits_to_indices(self, bits):\n bits = (bits >= 0).to(torch.long)\n indices = 2 ** torch.arange(\n 0,\n bits.shape[-1],\n 1,\n dtype=torch.long,\n device=bits.device,\n )\n return (bits * indices).sum(-1)\n\n def forward(self, z, half=False, collect_metrics=True):\n z = F.normalize(z, dim=-1)\n quantized, bsq_loss, metrics = self.bsq(z, collect_metrics=collect_metrics)\n if half:\n q_pre = quantized[:, :, :self.s1_bits]\n q_post = quantized[:, :, self.s1_bits:]\n z_indices = [self.bits_to_indices(q_pre), self.bits_to_indices(q_post)]\n else:\n z_indices = self.bits_to_indices(quantized)\n return bsq_loss, quantized, z_indices\n\n\nclass RMSNorm(torch.nn.Module):\n def __init__(self, dim: int, eps: float = 1e-5):\n super().__init__()\n self.eps = eps\n self.weight = nn.Parameter(torch.ones(dim))\n\n def _norm(self, x):\n return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)\n\n def forward(self, x):\n output = self._norm(x.float()).type_as(x)\n return output * self.weight\n\n\nclass FeedForward(nn.Module):\n def __init__(self, d_model, ff_dim, ffn_dropout_p=0.0):\n super().__init__()\n\n self.w1 = nn.Linear(d_model, ff_dim, bias=False)\n self.w3 = nn.Linear(d_model, ff_dim, bias=False)\n self.w2 = nn.Linear(ff_dim, d_model, bias=False)\n self.ffn_dropout = nn.Dropout(ffn_dropout_p)\n\n def forward(self, x):\n return self.ffn_dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))\n\n\nclass RotaryPositionalEmbedding(nn.Module):\n def __init__(self, dim):\n super().__init__()\n inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))\n self.register_buffer(\"inv_freq\", inv_freq)\n self.seq_len_cached = None\n self.cos_cached = None\n self.sin_cached = None\n\n def _update_cos_sin_cache(self, x, seq_len):\n if seq_len != self.seq_len_cached:\n self.seq_len_cached = seq_len\n t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)\n freqs = torch.einsum('i,j->ij', t, self.inv_freq)\n emb = torch.cat((freqs, freqs), dim=-1).to(x.device)\n self.cos_cached = emb.cos()[None, None, :, :]\n self.sin_cached = emb.sin()[None, None, :, :]\n return self.cos_cached, self.sin_cached\n\n def forward(self, q, k):\n cos, sin = self._update_cos_sin_cache(q, q.shape[-2])\n return (\n (q * cos) + (self._rotate_half(q) * sin),\n (k * cos) + (self._rotate_half(k) * sin),\n )\n\n def _rotate_half(self, x):\n x1, x2 = x.chunk(2, dim=-1)\n return torch.cat((-x2, x1), dim=-1)\n\n\nclass MultiHeadAttentionWithRoPE(nn.Module):\n def __init__(self, d_model, n_heads, attn_dropout_p=0.0, resid_dropout_p=0.0):\n super().__init__()\n self.d_model = d_model\n self.n_heads = n_heads\n self.head_dim = d_model // n_heads\n\n self.q_proj = nn.Linear(d_model, d_model)\n self.k_proj = nn.Linear(d_model, d_model)\n self.v_proj = nn.Linear(d_model, d_model)\n self.out_proj = nn.Linear(d_model, d_model)\n self.rotary = RotaryPositionalEmbedding(self.head_dim)\n self.attn_dropout_p = attn_dropout_p\n self.resid_dropout = nn.Dropout(resid_dropout_p)\n\n def forward(self, x, key_padding_mask=None):\n batch_size, seq_len, _ = x.shape\n\n q = self.q_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)\n k = self.k_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)\n v = self.v_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)\n\n q, k = self.rotary(q, k)\n\n if key_padding_mask is not None:\n attn_mask = key_padding_mask.unsqueeze(1).unsqueeze(2) # [batch, 1, 1, seq_len]\n attn_mask = attn_mask.expand(-1, self.n_heads, seq_len, -1) # [batch, n_heads, q_len, k_len]\n else:\n attn_mask = None\n\n attn_output = F.scaled_dot_product_attention(\n q, k, v,\n attn_mask=attn_mask,\n dropout_p=self.attn_dropout_p if self.training else 0.0,\n is_causal=True\n )\n\n attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)\n return self.resid_dropout(self.out_proj(attn_output))\n\n\nclass MultiHeadCrossAttentionWithRoPE(nn.Module):\n def __init__(self, d_model, n_heads, attn_dropout_p=0.0, resid_dropout=0.0):\n super().__init__()\n self.d_model = d_model\n self.n_heads = n_heads\n self.head_dim = d_model // n_heads\n\n self.q_proj = nn.Linear(d_model, d_model)\n self.k_proj = nn.Linear(d_model, d_model)\n self.v_proj = nn.Linear(d_model, d_model)\n self.out_proj = nn.Linear(d_model, d_model)\n self.rotary = RotaryPositionalEmbedding(self.head_dim)\n self.attn_dropout_p = attn_dropout_p\n self.resid_dropout = nn.Dropout(resid_dropout)\n\n def forward(self, query, key, value, key_padding_mask=None):\n batch_size, q_len, _ = query.shape\n _, seq_len, _ = key.shape\n\n q = self.q_proj(query).view(batch_size, q_len, self.n_heads, self.head_dim).transpose(1, 2)\n k = self.k_proj(key).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)\n v = self.v_proj(value).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)\n\n q, k = self.rotary(q, k)\n\n if key_padding_mask is not None:\n attn_mask = key_padding_mask.unsqueeze(1).unsqueeze(2)\n attn_mask = attn_mask.expand(-1, self.n_heads, q_len, -1)\n else:\n attn_mask = None\n\n is_causal_flag = self.training\n\n attn_output = F.scaled_dot_product_attention(\n q, k, v,\n attn_mask=attn_mask,\n dropout_p=self.attn_dropout_p if self.training else 0.0,\n is_causal=is_causal_flag\n )\n\n attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, q_len, self.d_model)\n return self.resid_dropout(self.out_proj(attn_output))\n\n\nclass HierarchicalEmbedding(nn.Module):\n def __init__(self, s1_bits, s2_bits, d_model=256):\n super().__init__()\n self.s1_bits = s1_bits\n self.s2_bits = s2_bits\n\n vocab_s1 = 2 ** s1_bits\n vocab_s2 = 2 ** s2_bits\n\n self.emb_s1 = nn.Embedding(vocab_s1, d_model)\n self.emb_s2 = nn.Embedding(vocab_s2, d_model)\n self.d_model = d_model\n self.fusion_proj = nn.Linear(d_model * 2, d_model)\n\n nn.init.normal_(self.emb_s1.weight, mean=0, std=d_model ** -0.5)\n nn.init.normal_(self.emb_s2.weight, mean=0, std=d_model ** -0.5)\n\n def split_token(self, token_ids: torch.Tensor, s2_bits: int):\n \"\"\"Inputs:\n token_ids (torch.Tensor): Composite token IDs of shape [batch_size, seq_len] or [N], each in range [0, 2^(s1_bits + s2_bits) - 1].\n s2_bits (int): Number of low bits used for the fine token (s2).\n \"\"\"\n assert isinstance(s2_bits, int) and s2_bits > 0, \"s2_bits must be a positive integer\"\n\n t = token_ids.long()\n mask = (1 \u003c\u003c s2_bits) - 1\n s2_ids = t & mask # extract low bits\n s1_ids = t >> s2_bits # extract high bits\n return s1_ids, s2_ids\n\n def forward(self, token_ids):\n \"\"\"Inputs:\n token_ids:\n - tuple or list: (s1_ids, s2_ids), each of shape [batch_size, seq_len], or\n - torch.Tensor: composite token IDs of shape [batch_size, seq_len], which will be split into (s1_ids, s2_ids) internally.\n Output: [batch_size, seq_len, d_model]\n \"\"\"\n if isinstance(token_ids, tuple) or isinstance(token_ids, list):\n s1_ids, s2_ids = token_ids\n else:\n s1_ids, s2_ids = self.split_token(token_ids, self.s2_bits)\n s1_emb = self.emb_s1(s1_ids) * math.sqrt(self.d_model)\n s2_emb = self.emb_s2(s2_ids) * math.sqrt(self.d_model)\n return self.fusion_proj(torch.cat([s1_emb, s2_emb], dim=-1))\n\n\nclass DependencyAwareLayer(nn.Module):\n def __init__(self, d_model, n_heads=4, attn_dropout_p=0.0, resid_dropout=0.0):\n super().__init__()\n self.cross_attn = MultiHeadCrossAttentionWithRoPE(d_model, n_heads, attn_dropout_p, resid_dropout)\n self.norm = RMSNorm(d_model)\n\n def forward(self, hidden_states, sibling_embed, key_padding_mask=None):\n \"\"\"hidden_states: [batch, seq_len, d_model]\n sibling_embed: Embedding from another subtoken\n \"\"\"\n attn_out = self.cross_attn(\n query=sibling_embed,\n key=hidden_states,\n value=hidden_states,\n key_padding_mask=key_padding_mask\n )\n return self.norm(hidden_states + attn_out)\n\n\nclass TransformerBlock(nn.Module):\n def __init__(self, d_model, n_heads, ff_dim=1024, ffn_dropout_p=0.0, attn_dropout_p=0.0, resid_dropout_p=0.0):\n super().__init__()\n self.norm1 = RMSNorm(d_model)\n self.self_attn = MultiHeadAttentionWithRoPE(d_model, n_heads, attn_dropout_p, resid_dropout_p)\n self.norm2 = RMSNorm(d_model)\n self.ffn = FeedForward(d_model, ff_dim, ffn_dropout_p)\n\n def forward(self, x, key_padding_mask=None):\n residual = x\n x = self.norm1(x)\n attn_out = self.self_attn(x, key_padding_mask=key_padding_mask)\n x = residual + attn_out\n\n residual = x\n x = self.norm2(x)\n ffn_out = self.ffn(x)\n x = residual + ffn_out\n return x\n\n\nclass DualHead(nn.Module):\n def __init__(self, s1_bits, s2_bits, d_model):\n super().__init__()\n self.vocab_s1 = 2 ** s1_bits\n self.vocab_s2 = 2 ** s2_bits\n self.proj_s1 = nn.Linear(d_model, self.vocab_s1)\n self.proj_s2 = nn.Linear(d_model, self.vocab_s2)\n\n def compute_loss(self, s1_logits, s2_logits, s1_targets, s2_targets, padding_mask=None):\n if padding_mask is not None:\n valid_mask = (padding_mask == 0)\n s1_logits = s1_logits[valid_mask]\n s2_logits = s2_logits[valid_mask]\n s1_targets = s1_targets[valid_mask]\n s2_targets = s2_targets[valid_mask]\n ce_s1 = F.cross_entropy(s1_logits, s1_targets)\n ce_s2 = F.cross_entropy(s2_logits, s2_targets)\n else:\n ce_s1 = F.cross_entropy(s1_logits.reshape(-1, self.vocab_s1), s1_targets.reshape(-1))\n ce_s2 = F.cross_entropy(s2_logits.reshape(-1, self.vocab_s2), s2_targets.reshape(-1))\n ce_loss = (ce_s1 + ce_s2) / 2\n return ce_loss, ce_s1, ce_s2\n\n def forward(self, x):\n return self.proj_s1(x)\n\n def cond_forward(self, x2):\n return self.proj_s2(x2)\n\n\nclass FixedEmbedding(nn.Module):\n def __init__(self, c_in, d_model):\n super(FixedEmbedding, self).__init__()\n\n w = torch.zeros(c_in, d_model).float()\n w.require_grad = False\n\n position = torch.arange(0, c_in).float().unsqueeze(1)\n div_term = (torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)).exp()\n\n w[:, 0::2] = torch.sin(position * div_term)\n w[:, 1::2] = torch.cos(position * div_term)\n\n self.emb = nn.Embedding(c_in, d_model)\n self.emb.weight = nn.Parameter(w, requires_grad=False)\n\n def forward(self, x):\n return self.emb(x).detach()\n\n\nclass TemporalEmbedding(nn.Module):\n def __init__(self, d_model, learn_pe):\n super(TemporalEmbedding, self).__init__()\n\n minute_size = 60\n hour_size = 24\n weekday_size = 7\n day_size = 32\n month_size = 13\n\n Embed = FixedEmbedding if not learn_pe else nn.Embedding\n self.minute_embed = Embed(minute_size, d_model)\n self.hour_embed = Embed(hour_size, d_model)\n self.weekday_embed = Embed(weekday_size, d_model)\n self.day_embed = Embed(day_size, d_model)\n self.month_embed = Embed(month_size, d_model)\n\n def forward(self, x):\n x = x.long()\n\n minute_x = self.minute_embed(x[:, :, 0])\n hour_x = self.hour_embed(x[:, :, 1])\n weekday_x = self.weekday_embed(x[:, :, 2])\n day_x = self.day_embed(x[:, :, 3])\n month_x = self.month_embed(x[:, :, 4])\n\n return hour_x + weekday_x + day_x + month_x + minute_x","content_type":"text/x-python; charset=utf-8","language":"python","size":22847,"content_sha256":"91e8ff4ea1d77a95e39be79a217c360165e585d90475c7a8369c2833953ba181"},{"filename":"scripts/utils/predictor/training.py","content":"import os\nimport sys\nimport time\nimport torch\nimport torch.nn as nn\nimport pandas as pd\nimport numpy as np\nimport json\nimport random\nfrom loguru import logger\nfrom datetime import datetime, timedelta\nfrom sentence_transformers import SentenceTransformer\nfrom dotenv import load_dotenv\n\n# Load environment variables\nload_dotenv()\n\n# Setup paths\nKRONOS_DIR = os.path.dirname(os.path.abspath(__file__))\nSRC_DIR = os.path.dirname(os.path.dirname(KRONOS_DIR))\nif SRC_DIR not in sys.path:\n sys.path.insert(0, SRC_DIR)\n\nfrom ..kronos.model import Kronos, KronosTokenizer, KronosPredictor\nfrom ..database_manager import DatabaseManager\nfrom ..stock_tools import StockTools\nfrom ..search_tools import SearchTools\nfrom ..llm.factory import get_model\nfrom ..visualizer import VisualizerTools\nfrom ..schema.models import ForecastResult, KLinePoint\nfrom agno.agent import Agent\n\nclass AutoSynthesisTrainer:\n def __init__(self, news_dim=384):\n self.device = \"cuda\" if torch.cuda.is_available() else \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n self.db = DatabaseManager()\n self.tools = StockTools(self.db)\n self.searcher = SearchTools(self.db)\n # Try loading from local cache first to avoid network timeouts\n model_name = os.getenv('EMBEDDING_MODEL', 'sentence-transformers/all-MiniLM-L6-v2')\n try:\n logger.info(f\"🔄 Attempting to load {model_name} from local cache...\")\n self.embedder = SentenceTransformer(model_name, device=self.device, local_files_only=True)\n logger.success(\"✅ Model loaded from local cache.\")\n except Exception:\n logger.warning(\"⚠️ Local cache not found or incomplete. Attempting to download...\")\n self.embedder = SentenceTransformer(model_name, device=self.device)\n self.news_dim = news_dim\n \n # Try loading from local cache first to avoid network timeouts\n try:\n logger.info(\"🔄 Attempting to load Kronos and Tokenizer from local cache...\")\n self.tokenizer = KronosTokenizer.from_pretrained(\"NeoQuasar/Kronos-Tokenizer-base\", local_files_only=True).to(self.device)\n base_model = Kronos.from_pretrained(\"NeoQuasar/Kronos-base\", local_files_only=True)\n logger.success(\"✅ Kronos and Tokenizer loaded from local cache.\")\n except Exception:\n logger.warning(\"⚠️ Local Kronos/Tokenizer not found or incomplete. Attempting to download...\")\n self.tokenizer = KronosTokenizer.from_pretrained(\"NeoQuasar/Kronos-Tokenizer-base\").to(self.device)\n base_model = Kronos.from_pretrained(\"NeoQuasar/Kronos-base\")\n \n self.model = Kronos(\n base_model.s1_bits, base_model.s2_bits, base_model.n_layers, \n base_model.d_model, base_model.n_heads, base_model.ff_dim,\n base_model.ffn_dropout_p, base_model.attn_dropout_p,\n base_model.resid_dropout_p, base_model.token_dropout_p,\n base_model.learn_te, news_dim=self.news_dim\n ).to(self.device)\n self.model.load_state_dict(base_model.state_dict(), strict=False)\n \n # LLM for causality verification\n provider = os.getenv(\"LLM_PROVIDER\", \"ust\")\n model_id = os.getenv(\"LLM_MODEL\", \"Qwen\")\n self.llm_agent = Agent(model=get_model(provider, model_id))\n\n def discover_shocks(self, ticker_list, threshold=2.0, limit_per_stock=5, days=365, pred_len=5):\n \"\"\"1. Find days with significant price movements (Look back 1 year)\"\"\"\n shocks = []\n end_date = datetime.now().strftime(\"%Y-%m-%d\")\n start_date = (datetime.now() - timedelta(days=days)).strftime(\"%Y-%m-%d\")\n \n for ticker in ticker_list:\n df = self.tools.get_stock_price(ticker, start_date=start_date, end_date=end_date)\n if df.empty or len(df) \u003c 60:\n continue\n \n # Look for big moves\n moves = df[df['change_pct'].abs() > threshold].copy()\n if moves.empty: continue\n \n count = 0\n for idx, row in moves.iterrows():\n # Ensure we have history before this day AND enough future days for eval\n date_idx = df.index.get_loc(idx)\n if date_idx \u003c 50 or date_idx + pred_len > len(df): continue\n \n shocks.append({\n 'ticker': ticker,\n 'date': row['date'],\n 'change': row['change_pct'],\n 'history': df.iloc[date_idx-50:date_idx],\n 'target': df.iloc[date_idx:date_idx + pred_len] # Now capturing pred_len days\n })\n count += 1\n if count >= limit_per_stock: break\n \n logger.info(f\"✨ Discovered {len(shocks)} potential price shocks over the last {days} days.\")\n return shocks\n\n def find_reason_and_verify(self, shock):\n \"\"\"2. Search for reasons and verify causality using LLM\"\"\"\n ticker_info = self.db.get_stock_by_code(shock['ticker'])\n name = ticker_info['name'] if ticker_info else shock['ticker']\n date_str = shock['date']\n \n # Try multiple query variations and engines\n queries = [\n f\"{name} ({shock['ticker']}) {date_str} 为什么涨跌原因\",\n f\"{name} {date_str} 异动原因\",\n f\"{shock['ticker']} {date_str} 新闻\"\n ]\n \n search_results = []\n for query in queries:\n logger.info(f\"🔍 Searching for reason: {query}\")\n # Try alternate engines\n for engine in [\"baidu\"]:\n try:\n results = self.searcher.search_list(query, engine=engine, max_results=3, enrich=False)\n if results:\n search_results = results\n break\n except Exception as e:\n logger.warning(f\"Search failed for {query} on {engine}: {e}\")\n \n if search_results:\n break\n time.sleep(random.uniform(1.0, 2.0))\n \n if not search_results:\n logger.warning(f\"⚠️ No search results found for {name} on {date_str} after multiple attempts.\")\n return None\n \n context = \"\\n\".join([f\"- {r['title']}: {r.get('content', '')[:300]}\" for r in search_results])\n \n prompt = f\"\"\"\n 任务：判断以下新闻是否解释了该股票在 {date_str} 的 {shock['change']:.2f}% 价格变动。\n \n 股票：{name}\n 日期：{date_str}\n 变动：{shock['change']:.2f}%\n \n 搜索结果：\n {context}\n \n 要求：\n 1. 该新闻是否在该日期左右发生？\n 2. 该新闻是否能逻辑上解释这种大幅波动（如财报、利好政策、重组、大环境暴跌等）？\n 3. 如果是，请总结一段 100 字以内的“核心推动原因”。\n 4. 返回 JSON: {{\"is_causal\": true/false, \"summary\": \"原因摘要\"}}\n \"\"\"\n \n try:\n res = self.llm_agent.run(prompt)\n data = json.loads(res.content.replace('```json', '').replace('```', '').strip())\n if data.get('is_causal'):\n logger.success(f\"✅ Verified cause for {name} on {date_str}: {data['summary']}\")\n return data['summary']\n else:\n logger.warning(f\"❌ Verified cause for {name} on {date_str}: {data['summary']}\")\n return None\n except Exception as e:\n logger.warning(f\"Verification failed: {e}\")\n return None\n\n def save_model(self, path=None):\n \"\"\"Save the news_proj weights\"\"\"\n if path is None:\n save_dir = os.path.join(SRC_DIR, \"exports/models\")\n os.makedirs(save_dir, exist_ok=True)\n path = os.path.join(save_dir, f\"kronos_news_v1_{datetime.now().strftime('%Y%m%d_%H%M')}.pt\")\n \n # We only really need to save the news_proj part as it's the only one we train\n torch.save({\n 'news_proj_state_dict': self.model.news_proj.state_dict(),\n 'news_dim': self.news_dim,\n 'd_model': self.model.d_model\n }, path)\n logger.success(f\"💾 Model weights saved to {path}\")\n return path\n\n def run_synthesis_and_train(self, tickers, pred_len=5):\n # 1. Discovery\n shocks = self.discover_shocks(tickers, pred_len=pred_len)\n print(f'find {len(shocks)} shocks')\n \n # 2. News Association & Verification\n dataset = []\n max_news_items = 200 # Limit to 200 news items per session to avoid search bans\n \n logger.info(f\"🧬 Starting News Association for {len(shocks)} shocks (Max limit: {max_news_items})\")\n \n for i, shock in enumerate(shocks):\n if len(dataset) >= max_news_items:\n logger.info(\"Reached maximum news items limit for this session.\")\n break\n \n summary = self.find_reason_and_verify(shock)\n if summary:\n # 3. Embedding news\n emb = self.embedder.encode(summary)\n dataset.append({\n 'history': shock['history'],\n 'target': shock['target'],\n 'news_emb': emb,\n 'summary': summary\n })\n \n # Add delay after search with randomness to avoid being blocked\n if i \u003c len(shocks) - 1:\n delay = random.uniform(2.0, 4.0)\n time.sleep(delay)\n \n if not dataset:\n logger.error(\"❌ No verified news-price pairs found. Adjust threshold or check if news is available in that period.\")\n return\n\n # 4. Train/Val Split\n random.seed(42)\n random.shuffle(dataset)\n \n if len(dataset) \u003c 2:\n train_set = dataset\n val_set = []\n logger.warning(f\"⚠️ Only {len(dataset)} sample(s) found. Training on all, skipping validation.\")\n else:\n split_idx = max(1, int(len(dataset) * 0.8))\n if split_idx >= len(dataset):\n split_idx = len(dataset) - 1\n \n train_set = dataset[:split_idx]\n val_set = dataset[split_idx:]\n logger.info(f\"🏗️ Dataset Split: {len(train_set)} samples for training, {len(val_set)} for validation.\")\n\n if not train_set:\n logger.error(\"❌ No samples for training.\")\n return\n\n # 5. Training (Few-shot)\n optimizer = torch.optim.Adam(self.model.news_proj.parameters(), lr=1e-3)\n criterion = nn.CrossEntropyLoss()\n self.model.train()\n \n loss_history = []\n logger.info(f\"🚀 Training for 30 epochs...\")\n for epoch in range(30):\n total_loss = 0\n for item in train_set:\n optimizer.zero_grad()\n \n # Prep Data\n hist_df = item['history']\n # For training, we still focus on the immediate next point (teacher forcing)\n target_df = item['target'].iloc[:1]\n \n hist_raw = hist_df[['open', 'high', 'low', 'close', 'volume']].values.astype(np.float32)\n hist_raw = np.column_stack([hist_raw, hist_raw[:, 3] * hist_raw[:, 4]]) \n \n mean, std = hist_raw.mean(axis=0), hist_raw.std(axis=0) + 1e-5\n hist_norm = torch.from_numpy((hist_raw - mean) / std).unsqueeze(0).to(self.device)\n \n target_raw = target_df[['open', 'high', 'low', 'close', 'volume']].values.astype(np.float32)\n target_raw = np.column_stack([target_raw, target_raw[:, 3] * target_raw[:, 4]])\n target_norm = torch.from_numpy((target_raw - mean) / std).unsqueeze(0).to(self.device)\n \n with torch.no_grad():\n z_indices = self.tokenizer.encode(hist_norm, half=True)\n t_indices = self.tokenizer.encode(target_norm, half=True)\n s1_ids, s2_ids = z_indices[0], z_indices[1]\n t_s1, t_s2 = t_indices[0], t_indices[1]\n \n news_t = torch.from_numpy(item['news_emb']).unsqueeze(0).to(self.device)\n s1_logits, s2_logits = self.model(s1_ids, s2_ids, news_emb=news_t, use_teacher_forcing=True, s1_targets=t_s1)\n \n loss = (criterion(s1_logits[:, -1, :], t_s1[:, 0]) + criterion(s2_logits[:, -1, :], t_s2[:, 0])) / 2\n loss.backward()\n optimizer.step()\n total_loss += loss.item()\n \n avg_epoch_loss = total_loss / max(1, len(train_set))\n loss_history.append(avg_epoch_loss)\n \n if (epoch + 1) % 10 == 0:\n logger.info(f\"Epoch {epoch+1} Loss: {avg_epoch_loss:.4f}\")\n\n # 5.1 Visualize Loss Curve\n loss_chart = VisualizerTools.generate_loss_chart(loss_history)\n VisualizerTools.render_chart_to_file(loss_chart, os.path.join(SRC_DIR, \"exports/training_results/loss_curve.html\"))\n\n # 5.2 Save final model\n self.save_model()\n\n # 6. Final Evaluation on Validation Set\n if not val_set:\n logger.warning(\"⚠️ Validation set is empty. Skipping statistical analysis.\")\n return\n\n logger.info(f\"🧪 Final Evaluation: Base vs News-Integrated ({pred_len}-day Window)\")\n self.model.eval()\n predictor = KronosPredictor(self.model, self.tokenizer, device=self.device)\n \n base_maes = []\n news_maes = []\n \n print(\"\\n\" + \"=\"*90)\n print(f\"{'Date':\u003c12} | {'Ticker':\u003c8} | {'Base MAE':\u003c15} | {'News MAE':\u003c15} | {'Improvement'}\")\n print(\"-\" * 90)\n\n for item in val_set:\n h = item['history']\n t = item['target']\n actuals = t['close'].values[:pred_len]\n \n x_ts = pd.to_datetime(h['date'])\n # Future timestamps: handle business days if possible, or just simple offset\n future_dates = pd.date_range(start=x_ts.iloc[-1] + timedelta(days=1), periods=pred_len, freq='B')\n y_ts = pd.Series(future_dates)\n \n # A. Base Prediction\n p_base = predictor.predict(h, x_ts, y_ts, pred_len=pred_len, news_emb=None, verbose=False)\n b_preds = p_base['close'].values[:len(actuals)]\n \n # B. News-Aware Prediction\n p_news = predictor.predict(h, x_ts, y_ts, pred_len=pred_len, news_emb=item['news_emb'], verbose=False)\n n_preds = p_news['close'].values[:len(actuals)]\n \n # Calculate MAE over the window\n b_mae = np.mean(np.abs(b_preds - actuals))\n n_mae = np.mean(np.abs(n_preds - actuals))\n \n base_maes.append(b_mae)\n news_maes.append(n_mae)\n \n improvement = (b_mae - n_mae) / (b_mae + 1e-6) * 100\n \n date_str = str(t['date'].values[0])[:10]\n ticker = h.iloc[-1]['ticker'] if 'ticker' in h.columns else \"Stock\"\n print(f\"{date_str:\u003c12} | {ticker:\u003c8} | {b_mae:\u003c15.4f} | {n_mae:\u003c15.4f} | {improvement:>+7.1f}%\")\n\n # C. Generate Visualization for this case\n try:\n # Helper to convert DF to KLinePoints\n def to_kp_list(preds_df):\n points = []\n for idx, row in preds_df.iterrows():\n points.append(KLinePoint(\n date=str(idx)[:10],\n open=row['open'],\n high=row['high'],\n low=row['low'],\n close=row['close'],\n volume=row['volume'] if 'volume' in row else 0\n ))\n return points\n\n forecast_obj = ForecastResult(\n ticker=ticker,\n base_forecast=to_kp_list(p_base),\n adjusted_forecast=to_kp_list(p_news),\n rationale=item['summary']\n )\n\n # Ground truth for visualizer expects a DataFrame with 'date' and 'close'\n gt_df = t[['date', 'open', 'high', 'low', 'close', 'volume']]\n \n chart = VisualizerTools.generate_stock_chart(\n df=h, \n ticker=ticker, \n title=f\"Training Eval: {ticker} ({date_str}) Improvement: {improvement:.1f}%\",\n forecast=forecast_obj,\n ground_truth=gt_df\n )\n \n safe_date = date_str.replace(\"-\", \"\")\n filename = f\"eval_{ticker}_{safe_date}.html\"\n VisualizerTools.render_chart_to_file(chart, os.path.join(SRC_DIR, f\"exports/training_results/{filename}\"))\n except Exception as e:\n logger.error(f\"Failed to generate eval chart for {ticker}: {e}\")\n\n # Summary Statistics\n avg_base_err = sum(base_maes) / max(1, len(base_maes))\n avg_news_err = sum(news_maes) / max(1, len(news_maes))\n overall_imp = (avg_base_err - avg_news_err) / (avg_base_err + 1e-6) * 100\n \n print(\"-\" * 90)\n print(f\"{'AVERAGE':\u003c12} | {'-':\u003c8} | {avg_base_err:\u003c15.4f} | {avg_news_err:\u003c15.4f} | {overall_imp:>+7.1f}%\")\n print(\"=\"*90 + \"\\n\")\n \n logger.success(f\"🏁 Statistical Analysis Complete. Avg Error Reduction ({pred_len}-day): {overall_imp:.2f}%\")\n logger.info(f\"📊 Visualization results saved to: {os.path.join(SRC_DIR, 'exports/training_results/')}\")\n\nif __name__ == \"__main__\":\n trainer = AutoSynthesisTrainer()\n \n logger.info(\"📂 Fetching all stock codes from database...\")\n res = trainer.db.execute_query(\"SELECT code FROM stock_list\")\n all_tickers = [row['code'] for row in res]\n \n if not all_tickers:\n logger.warning(\"⚠️ No tickers found in stock_list table. Trying to sync...\")\n trainer.tools._check_and_update_stock_list(force=True)\n res = trainer.db.execute_query(\"SELECT code FROM stock_list\")\n all_tickers = [row['code'] for row in res]\n\n logger.info(f\"🚀 Starting training on potential stocks (1-year scan)...\")\n # 为了演示，我们扫描前 100 个股票，寻找最近一年的冲击点\n trainer.run_synthesis_and_train(all_tickers[:100], pred_len=1)\n","content_type":"text/x-python; charset=utf-8","language":"python","size":18696,"content_sha256":"8cbc820f3142abce635dc9b8baf680434a3947219476cc5d430ed799549905a9"},{"filename":"scripts/utils/search_tools.py","content":"import os\nimport hashlib\nimport json\nimport re\nimport requests\nimport time\nimport threading\nfrom typing import List, Dict, Optional, Any\nfrom agno.tools.duckduckgo import DuckDuckGoTools\nfrom agno.tools.baidusearch import BaiduSearchTools\nfrom agno.agent import Agent\nfrom loguru import logger\nfrom datetime import datetime\nfrom .database_manager import DatabaseManager\nfrom .content_extractor import ContentExtractor\nfrom .llm.factory import get_model\nfrom .hybrid_search import LocalNewsSearch\n\n# 默认搜索缓存 TTL（秒），可通过环境变量覆盖\nDEFAULT_SEARCH_TTL = int(os.getenv(\"SEARCH_CACHE_TTL\", \"3600\")) # 默认 1 小时\n\n\nclass JinaSearchEngine:\n \"\"\"Jina Search API 封装 - 使用 s.jina.ai 进行网络搜索\"\"\"\n \n JINA_SEARCH_URL = \"https://s.jina.ai/\"\n \n # 速率限制配置\n _rate_limit_no_key = 10 # 无 key 时每分钟最大请求数\n _rate_window = 60.0\n _min_interval = 2.0\n _request_times = []\n _last_request_time = 0.0\n _lock = threading.Lock()\n \n def __init__(self):\n self.api_key = os.getenv(\"JINA_API_KEY\", \"\").strip()\n self.has_api_key = bool(self.api_key)\n if self.has_api_key:\n logger.info(\"✅ Jina Search API key configured\")\n \n @classmethod\n def _wait_for_rate_limit(cls, has_api_key: bool) -> None:\n \"\"\"等待以满足速率限制\"\"\"\n if has_api_key:\n time.sleep(0.3)\n return\n \n with cls._lock:\n current_time = time.time()\n cls._request_times = [t for t in cls._request_times if current_time - t \u003c cls._rate_window]\n \n if len(cls._request_times) >= cls._rate_limit_no_key:\n oldest = cls._request_times[0]\n wait_time = cls._rate_window - (current_time - oldest) + 1.0\n if wait_time > 0:\n logger.warning(f\"⏳ Jina Search rate limit, waiting {wait_time:.1f}s...\")\n time.sleep(wait_time)\n current_time = time.time()\n cls._request_times = [t for t in cls._request_times if current_time - t \u003c cls._rate_window]\n \n time_since_last = current_time - cls._last_request_time\n if time_since_last \u003c cls._min_interval:\n time.sleep(cls._min_interval - time_since_last)\n \n cls._request_times.append(time.time())\n cls._last_request_time = time.time()\n \n def search(self, query: str, max_results: int = 5) -> List[Dict]:\n \"\"\"\n 使用 Jina Search API 执行搜索\n \n Args:\n query: 搜索关键词\n max_results: 返回结果数量\n \n Returns:\n 搜索结果列表，每个结果包含 title, url, content\n \"\"\"\n if not query:\n return []\n \n logger.info(f\"🔍 Jina Search: {query}\")\n \n # 等待速率限制\n self._wait_for_rate_limit(self.has_api_key)\n \n headers = {\n \"Accept\": \"application/json\",\n \"X-Retain-Images\": \"none\",\n }\n \n if self.has_api_key:\n headers[\"Authorization\"] = f\"Bearer {self.api_key}\"\n \n try:\n # Jina Search API: https://s.jina.ai/{query}\n import urllib.parse\n encoded_query = urllib.parse.quote(query)\n url = f\"{self.JINA_SEARCH_URL}{encoded_query}\"\n \n response = requests.get(url, headers=headers, timeout=30)\n \n if response.status_code == 429:\n logger.warning(\"⚠️ Jina Search rate limited (429), waiting 30s...\")\n time.sleep(30)\n return self.search(query, max_results)\n \n if response.status_code != 200:\n logger.warning(f\"Jina Search failed (Status {response.status_code})\")\n return []\n \n # 解析响应\n try:\n data = response.json()\n except json.JSONDecodeError:\n # 如果返回纯文本，尝试解析\n data = {\"data\": [{\"title\": \"Search Result\", \"url\": \"\", \"content\": response.text}]}\n \n results = []\n \n # Jina 返回格式可能是 {\"data\": [...]} 或直接是列表\n items = data.get(\"data\", []) if isinstance(data, dict) else data\n if not isinstance(items, list):\n items = [items] if items else []\n \n for i, item in enumerate(items[:max_results]):\n if isinstance(item, dict):\n results.append({\n \"title\": item.get(\"title\", f\"Result {i+1}\"),\n \"url\": item.get(\"url\", \"\"),\n \"href\": item.get(\"url\", \"\"), # 兼容性\n \"content\": item.get(\"content\", item.get(\"description\", \"\")),\n \"body\": item.get(\"content\", item.get(\"description\", \"\")), # 兼容性\n })\n elif isinstance(item, str):\n results.append({\n \"title\": f\"Result {i+1}\",\n \"url\": \"\",\n \"content\": item\n })\n \n logger.info(f\"✅ Jina Search returned {len(results)} results\")\n return results\n \n except requests.exceptions.Timeout:\n logger.error(\"Jina Search timeout\")\n return []\n except requests.exceptions.RequestException as e:\n logger.error(f\"Jina Search request error: {e}\")\n return []\n except Exception as e:\n logger.error(f\"Jina Search unexpected error: {e}\")\n return []\n\nclass SearchTools:\n \"\"\"扩展性搜索工具库 - 支持多引擎聚合与内容缓存\"\"\"\n \n def __init__(self, db: DatabaseManager):\n self.db = db\n \n # 检查 Jina API Key 是否配置\n jina_api_key = os.getenv(\"JINA_API_KEY\", \"\").strip()\n self._jina_enabled = bool(jina_api_key)\n \n self._engines = {\n \"ddg\": DuckDuckGoTools(),\n \"baidu\": BaiduSearchTools(),\n \"local\": LocalNewsSearch(db)\n }\n \n # 如果配置了 Jina API Key，添加 Jina 引擎\n if self._jina_enabled:\n self._engines[\"jina\"] = JinaSearchEngine()\n logger.info(\"🚀 Jina Search engine enabled (JINA_API_KEY configured)\")\n \n # 确定默认搜索引擎\n self._default_engine = \"jina\" if self._jina_enabled else \"ddg\"\n\n def _generate_hash(self, query: str, engine: str, max_results: int) -> str:\n return hashlib.md5(f\"{engine}:{query}:{max_results}\".encode()).hexdigest()\n\n def search(self, query: str, engine: str = None, max_results: int = 5, ttl: Optional[int] = None) -> str:\n \"\"\"\n 使用指定搜索引擎执行网络搜索，结果会被缓存以提高效率。\n \n Args:\n query: 搜索关键词，如 \"英伟达财报\" 或 \"光伏行业政策\"。\n engine: 搜索引擎选择。可选值: \n \"jina\" (Jina Search，需配置 JINA_API_KEY，LLM友好输出),\n \"ddg\" (DuckDuckGo，推荐英文/国际搜索), \n \"baidu\" (百度，推荐中文/国内搜索),\n \"local\" (本地历史新闻搜索，基于向量+BM25)。\n 默认: 若配置了 JINA_API_KEY 则使用 \"jina\"，否则 \"ddg\"。\n max_results: 期望返回的结果数量，默认 5 条。\n ttl: 缓存有效期（秒）。如果缓存超过此时间会重新搜索。\n 默认使用环境变量 SEARCH_CACHE_TTL 或 3600 秒。\n 设为 0 可强制刷新。\n \n Returns:\n 搜索结果的文本描述，包含标题、摘要和链接。\n \"\"\"\n # 使用默认引擎（如果配置了 Jina 则优先使用 Jina）\n if engine is None:\n engine = self._default_engine\n \n if engine not in self._engines:\n return f\"Error: Unsupported engine '{engine}'. Available: {list(self._engines.keys())}\"\n\n query_hash = self._generate_hash(query, engine, max_results)\n effective_ttl = ttl if ttl is not None else DEFAULT_SEARCH_TTL\n \n # 1. 尝试从缓存读取 (local 引擎不缓存，因为它本身就是查库)\n if engine != \"local\":\n cache = self.db.get_search_cache(query_hash, ttl_seconds=effective_ttl if effective_ttl > 0 else None)\n if cache and effective_ttl != 0:\n logger.info(f\"ℹ️ Found search results in cache for: {query} ({engine})\")\n return cache['results']\n\n # 2. 执行真实搜索\n logger.info(f\"📡 Searching {engine} for: {query}\")\n try:\n tool = self._engines[engine]\n if engine == \"jina\":\n # Jina Search 返回 List[Dict]\n jina_results = tool.search(query, max_results=max_results)\n results = []\n for r in jina_results:\n results.append({\n \"title\": r.get(\"title\", \"\"),\n \"href\": r.get(\"url\", \"\"),\n \"body\": r.get(\"content\", \"\")\n })\n elif engine == \"ddg\":\n results = tool.duckduckgo_search(query, max_results=max_results)\n elif engine == \"baidu\":\n results = tool.baidu_search(query, max_results=max_results)\n elif engine == \"local\":\n # LocalNewsSearch 返回的是 List[Dict]\n local_results = tool.search(query, top_n=max_results)\n results = []\n for r in local_results:\n results.append({\n \"title\": r.get(\"title\"),\n \"href\": r.get(\"url\", \"local\"),\n \"body\": r.get(\"content\", \"\")\n })\n else:\n results = \"Search not implemented for this engine.\"\n \n results_str = str(results)\n if engine != \"local\":\n self.db.save_search_cache(query_hash, query, engine, results_str)\n return results_str\n \n except Exception as e:\n # 搜索失败时的降级策略\n if engine == \"jina\":\n logger.warning(f\"⚠️ Jina search failed, falling back to ddg: {query} ({e})\")\n try:\n return self.search(query, engine=\"ddg\", max_results=max_results, ttl=ttl)\n except Exception as e2:\n logger.error(f\"❌ DDG fallback also failed for {query}: {e2}\")\n elif engine == \"ddg\":\n logger.warning(f\"⚠️ DDG search failed, falling back to baidu: {query} ({e})\")\n try:\n return self.search(query, engine=\"baidu\", max_results=max_results, ttl=ttl)\n except Exception as e2:\n logger.error(f\"❌ Baidu fallback also failed for {query}: {e2}\")\n\n logger.error(f\"❌ Search failed for {query}: {e}\")\n return f\"Error occurred during search: {str(e)}\"\n\n def search_list(self, query: str, engine: str = None, max_results: int = 5, ttl: Optional[int] = None, enrich: bool = True) -> List[Dict]:\n \"\"\"\n 执行搜索并返回结构化列表 (List[Dict])。\n Dict 包含: title, href (or url), body (or snippet)\n \n Args:\n engine: 搜索引擎，默认使用配置的默认引擎（Jina 优先）\n enrich: 是否抓取正文内容 (默认 True)\n \"\"\"\n # 使用默认引擎\n if engine is None:\n engine = self._default_engine\n \n if engine not in self._engines:\n logger.error(f\"Unsupported engine {engine}\")\n return []\n \n # 不同的 hash 以区分是否 enrichment\n enrich_suffix = \":enriched\" if enrich else \"\"\n query_hash = self._generate_hash(query, engine + enrich_suffix, max_results)\n effective_ttl = ttl if ttl is not None else DEFAULT_SEARCH_TTL\n \n # 1. 尝试从缓存读取\n cache = self.db.get_search_cache(query_hash, ttl_seconds=effective_ttl if effective_ttl > 0 else None)\n if cache and effective_ttl != 0:\n try:\n cached_data = json.loads(cache['results'])\n if isinstance(cached_data, list):\n logger.info(f\"ℹ️ Found structured search cache for: {query}\")\n return cached_data\n except:\n pass\n \n # 1.5 Smart Cache (Fuzzy + LLM)\n if effective_ttl != 0:\n try:\n # 1. Similar cached queries\n similar_queries = self.db.find_similar_queries(query, limit=3)\n # Filter by TTL\n valid_candidates = []\n for q in similar_queries:\n if q['query'] == query: continue \n q_time = datetime.fromisoformat(q['timestamp'])\n if effective_ttl and (datetime.now() - q_time).total_seconds() > effective_ttl:\n continue\n q['type'] = 'cached_search'\n valid_candidates.append(q)\n\n # 2. Relevant local news (as search results)\n local_news = self.db.search_local_news(query, limit=3)\n if local_news:\n # Group local news as a single \"candidate\" source? Or individual?\n # Better to treat \"Local News Database\" as one candidate source that contains X items.\n # Or just add them to candidates list?\n # Let's package strictly relevant news as a \"local_news_bundle\"\n valid_candidates.append({\n 'type': 'local_news',\n 'query': 'Local Database News',\n 'items': local_news,\n 'timestamp': datetime.now().isoformat()\n })\n \n if valid_candidates:\n logger.info(f\"🤔 Found {len(valid_candidates)} smart cache candidates (Queries/News). Asking LLM...\")\n evaluation = self._evaluate_cache_relevance(query, valid_candidates)\n \n if evaluation and evaluation.get('reuse', False):\n idx = evaluation.get('index', -1)\n if 0 \u003c= idx \u003c len(valid_candidates):\n chosen = valid_candidates[idx]\n logger.info(f\"🤖 LLM suggested reusing: '{chosen.get('query')}' ({chosen['type']})\")\n \n if chosen['type'] == 'cached_search':\n # Load the chosen cache\n cache = self.db.get_search_cache(chosen['query_hash']) \n if cache:\n try:\n cached_data = json.loads(cache['results'])\n if isinstance(cached_data, list):\n return cached_data\n except:\n pass\n elif chosen['type'] == 'local_news':\n # Convert local news items to search result format\n news_results = []\n for i, news in enumerate(chosen['items'], 1):\n news_results.append({\n \"id\": news.get('id'),\n \"rank\": i,\n \"title\": news.get('title'),\n \"url\": news.get('url'),\n \"content\": news.get('content'),\n \"original_snippet\": news.get('content')[:200] if news.get('content') else '',\n \"source\": f\"Local News ({news.get('source')})\",\n \"publish_time\": news.get('publish_time'),\n \"crawl_time\": news.get('crawl_time'),\n \"sentiment_score\": news.get('sentiment_score', 0),\n \"meta_data\": {\"origin\": \"local_db\"}\n })\n return news_results\n\n except Exception as e:\n logger.warning(f\"Smart cache check failed: {e}\")\n \n # 2. 执行搜索\n logger.info(f\"📡 Searching {engine} (structured) for: {query}\")\n try:\n tool = self._engines[engine]\n results = []\n if engine == \"jina\":\n # Jina Search 直接返回结构化数据\n jina_results = tool.search(query, max_results=max_results)\n for r in jina_results:\n results.append({\n \"title\": r.get(\"title\", \"\"),\n \"url\": r.get(\"url\", \"\"),\n \"href\": r.get(\"url\", \"\"),\n \"body\": r.get(\"content\", \"\"),\n \"content\": r.get(\"content\", \"\"),\n \"source\": \"Jina Search\"\n })\n elif engine == \"ddg\":\n results = tool.duckduckgo_search(query, max_results=max_results)\n elif engine == \"baidu\":\n results = tool.baidu_search(query, max_results=max_results)\n elif engine == \"local\":\n # LocalNewsSearch 返回的是 List[Dict]\n local_results = tool.search(query, top_n=max_results)\n results = []\n for r in local_results:\n results.append({\n \"title\": r.get(\"title\"),\n \"url\": r.get(\"url\", \"local\"),\n \"body\": r.get(\"content\", \"\")[:500],\n \"source\": f\"Local ({r.get('source', 'db')})\",\n \"publish_time\": r.get(\"publish_time\")\n })\n \n # 处理字符串类型的 JSON 返回 (Baidu 常返 JSON 字符串)\n if isinstance(results, str) and engine not in [\"local\", \"jina\"]:\n try:\n results = json.loads(results)\n except:\n pass\n \n # 转为统一格式\n normalized_results = []\n if isinstance(results, list):\n \n for i, r in enumerate(results, 1):\n title = r.get('title', '')\n url = r.get('href') or r.get('url') or r.get('link', '')\n content = r.get('body') or r.get('snippet') or r.get('abstract', '')\n \n if title and url:\n normalized_results.append({\n \"id\": self._generate_hash(url + query, \"search_item\", i),\n \"rank\": i,\n \"title\": title,\n \"url\": url,\n \"content\": content,\n \"original_snippet\": content, # 保留摘要\n \"source\": f\"Search ({engine})\",\n \"publish_time\": datetime.now().isoformat(), # 暂用当前时间\n \"crawl_time\": datetime.now().isoformat(),\n \"meta_data\": {\"query\": query, \"engine\": engine}\n })\n \n # Fallback if still string and failed to parse\n elif isinstance(results, str) and results:\n normalized_results.append({\"title\": query, \"url\": \"\", \"content\": results, \"source\": engine})\n\n # 3. 抓取正文 & 计算情绪 (Enrichment)\n # 注意：如果使用 Jina Search，内容已经是 LLM 友好格式，可选择跳过 enrichment\n skip_content_enrichment = (engine == \"jina\")\n \n if enrich and normalized_results:\n logger.info(f\"🕸️ Enriching {len(normalized_results)} search results with Jina & Sentiment...\")\n extractor = ContentExtractor()\n \n # Lazy load sentiment tool\n if not hasattr(self, 'sentiment_tool') or self.sentiment_tool is None:\n from ..sentiment_tools import SentimentTools\n self.sentiment_tool = SentimentTools(self.db)\n \n for item in normalized_results:\n if item.get(\"url\"):\n try:\n # 如果是 Jina Search，内容已经足够好，跳过额外抓取\n if skip_content_enrichment and item.get(\"content\") and len(item.get(\"content\", \"\")) > 100:\n full_content = item[\"content\"]\n else:\n # Use Jina Reader to get full content\n full_content = extractor.extract_with_jina(item[\"url\"], timeout=60)\n \n if full_content and len(full_content) > 100:\n item[\"content\"] = full_content\n \n # Calculate sentiment\n # Use title + snippet of content for efficiency\n text_to_analyze = f\"{item['title']} {full_content[:500]}\"\n sent_result = self.sentiment_tool.analyze_sentiment(text_to_analyze) # Using self.sentiment_tool\n score = sent_result.get('score', 0.0)\n item[\"sentiment_score\"] = float(score)\n \n logger.info(f\" ✅ Enriched: {item['title'][:20]}... (Sentiment: {score:.2f})\")\n else:\n # Fallback: Use snippet for sentiment\n logger.info(f\" ⚠️ Content short/failed for {item['url']}, using snippet for sentiment.\")\n text_to_analyze = f\"{item['title']} {item['content']}\" # content is snippet here\n sent_result = self.sentiment_tool.analyze_sentiment(text_to_analyze)\n score = sent_result.get('score', 0.0)\n item[\"sentiment_score\"] = float(score)\n\n except Exception as e:\n # Fallback: Use snippet for sentiment on error\n logger.warning(f\"Failed to enrich {item['url']}: {e}. Using snippet.\")\n text_to_analyze = f\"{item['title']} {item['content']}\"\n sent_result = self.sentiment_tool.analyze_sentiment(text_to_analyze)\n score = sent_result.get('score', 0.0)\n item[\"sentiment_score\"] = float(score)\n \n # 缓存结果 list\n if normalized_results:\n # Pass list directly, DB manager will handle JSON dump for main cache and populate search_details\n # Only cache if NOT from local news reuse (though this logic path is for fresh search)\n self.db.save_search_cache(query_hash, query, engine, normalized_results)\n \n return normalized_results\n \n except Exception as e:\n # 搜索失败时的降级策略\n if engine == \"jina\":\n logger.warning(f\"⚠️ Jina search_list failed, falling back to ddg: {query} ({e})\")\n try:\n return self.search_list(query, engine=\"ddg\", max_results=max_results, ttl=ttl, enrich=enrich)\n except Exception as e2:\n logger.error(f\"❌ DDG fallback (search_list) also failed for {query}: {e2}\")\n elif engine == \"ddg\":\n logger.warning(f\"⚠️ DDG search_list failed, falling back to baidu: {query} ({e})\")\n try:\n return self.search_list(query, engine=\"baidu\", max_results=max_results, ttl=ttl, enrich=enrich)\n except Exception as e2:\n logger.error(f\"❌ Baidu fallback (search_list) also failed for {query}: {e2}\")\n\n logger.error(f\"❌ Structured search failed for {query}: {e}\")\n return []\n\n def _evaluate_cache_relevance(self, current_query: str, candidates: List[Dict]) -> Dict:\n \"\"\"\n 使用 LLM 评估缓存候选是否足以回答当前问题。\n \"\"\"\n try:\n # Prepare candidates text\n candidates_desc = []\n for i, c in enumerate(candidates):\n if c['type'] == 'cached_search':\n # Preview cached results if available? \n # Maybe just use the query string as a proxy for what's in there.\n # Or peek at 'results' snippet.\n preview = \"\"\n try:\n # Attempt to peek first result title from JSON string\n # Note: c.get('results') might be a stringified JSON list\n res_list = json.loads(c.get('results', '[]'))\n if res_list and isinstance(res_list, list) and len(res_list) > 0:\n first_item = res_list[0]\n if isinstance(first_item, dict) and 'title' in first_item:\n preview = f\" (Contains: {first_item.get('title', '')[:50]}...)\"\n except:\n pass\n candidates_desc.append(f\"[{i}] Old Search Query: '{c['query']}' {preview} (Time: {c['timestamp']})\")\n elif c['type'] == 'local_news':\n # List titles of local news\n titles = [item['title'] for item in c['items'][:3]]\n candidates_desc.append(f\"[{i}] Local Database News: {', '.join(titles)}... (Time: {c['timestamp']})\")\n\n prompt = f\"\"\"\n Task: Decide if existing information is sufficient for the new search query.\n \n New Query: \"{current_query}\"\n \n Available Information Candidates:\n {chr(10).join(candidates_desc)}\n \n Instructions:\n 1. Analyze if any candidate provides ENOUGH up-to-date info for the \"New Query\".\n 2. If yes, choose the best one.\n 3. If the query implies needing LATEST real-time info and candidates are old, choose none.\n 4. Return strictly JSON: {{\"reuse\": true/false, \"index\": \u003ccandidate_index_int>, \"reason\": \"short explanation\"}}\n \"\"\"\n # 初始化模型\n provider = os.getenv(\"LLM_PROVIDER\", \"ust\")\n model_id = os.getenv(\"LLM_MODEL\", \"Qwen\")\n host = os.getenv(\"LLM_HOST\")\n if host:\n model = get_model(provider, model_id, host=host)\n else:\n model = get_model(provider, model_id)\n \n agent = Agent(model=model, markdown=True)\n \n response = agent.run(prompt)\n content = response.content\n \n # Parse JSON\n json_match = re.search(r'```json\\s*(.*?)\\s*```', content, re.DOTALL)\n if json_match:\n return json.loads(json_match.group(1))\n elif '{' in content:\n # Fallback for cases where LLM doesn't wrap in ```json\n return json.loads(content[content.find('{'):content.rfind('}')+1])\n return {\"reuse\": False}\n \n except Exception as e:\n logger.warning(f\"LLM evaluation failed: {e}\")\n return {\"reuse\": False}\n\n def aggregate_search(self, query: str, engines: Optional[List[str]] = None, max_results: int = 5) -> str:\n \"\"\"\n 使用多个搜索引擎同时搜索并聚合结果，获得更全面的信息覆盖。\n \n Args:\n query: 搜索关键词。\n engines: 要使用的搜索引擎列表。可选值: [\"ddg\", \"baidu\"]。\n 默认同时使用 ddg 和 baidu。\n max_results: 每个引擎期望返回的结果数量。\n \n Returns:\n 聚合后的搜索结果，按引擎分组显示。\n \"\"\"\n engines = engines or [\"ddg\", \"baidu\"]\n aggregated_results = []\n for engine in engines:\n res = self.search(query, engine=engine, max_results=max_results)\n aggregated_results.append(f\"--- Results from {engine.upper()} ---\\n{res}\")\n \n return \"\\n\\n\".join(aggregated_results)\n","content_type":"text/x-python; charset=utf-8","language":"python","size":29258,"content_sha256":"839c578de49ee139f9c7a4d36cdcafddc833ef407f73fdb0e38e0c02abac6ad9"},{"filename":"scripts/utils/sentiment_tools.py","content":"import os\nfrom typing import Dict, List, Union, Optional\nimport json\nfrom loguru import logger\nfrom agno.agent import Agent\nfrom .llm.factory import get_model\nfrom .database_manager import DatabaseManager\n\n# 从环境变量读取默认情绪分析模式\nDEFAULT_SENTIMENT_MODE = os.getenv(\"SENTIMENT_MODE\", \"auto\") # auto, bert, llm\n\nclass SentimentTools:\n \"\"\"\n 情绪分析工具 - 支持 LLM 和 BERT 两种模式\n \n 模式说明:\n - \"auto\": 自动选择，优先使用 BERT（速度快），不可用时回退到 LLM\n - \"bert\": 强制使用 BERT 模型（需要 transformers 库）\n - \"llm\": 强制使用 LLM（更准确但较慢）\n \n 可通过环境变量 SENTIMENT_MODE 设置默认模式。\n \"\"\"\n \n def __init__(self, db: DatabaseManager, mode: Optional[str] = None, \n model_provider: str = \"openai\", model_id: str = \"gpt-4o\"):\n \"\"\"\n 初始化情绪分析工具。\n \n Args:\n db: 数据库管理器实例\n mode: 分析模式，可选 \"auto\", \"bert\", \"llm\"。None 则使用环境变量默认值。\n model_provider: LLM 提供商，如 \"openai\", \"ust\", \"deepseek\"\n model_id: 模型标识符\n \"\"\"\n self.db = db\n self.mode = mode or DEFAULT_SENTIMENT_MODE\n self.llm_model = None\n self.bert_pipeline = None\n \n # Initialize LLM\n try:\n provider = \"ust\" if os.getenv(\"UST_KEY_API\") else model_provider\n m_id = \"Qwen\" if provider == \"ust\" else model_id\n self.llm_model = get_model(provider, m_id)\n except Exception as e:\n logger.warning(f\"LLM initialization skipped: {e}\")\n\n # Initialize BERT if needed\n if self.mode in [\"bert\", \"auto\"]:\n try:\n from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification\n from transformers.utils import logging as transformers_logging\n transformers_logging.set_verbosity_error() # 减少冗余日志\n \n bert_model = os.getenv(\"BERT_SENTIMENT_MODEL\", \"uer/roberta-base-finetuned-chinanews-chinese\")\n \n # 优先使用本地缓存\n try:\n tokenizer = AutoTokenizer.from_pretrained(bert_model, local_files_only=True)\n model = AutoModelForSequenceClassification.from_pretrained(bert_model, local_files_only=True)\n \n self.bert_pipeline = pipeline(\n \"sentiment-analysis\", \n model=model,\n tokenizer=tokenizer,\n device=-1\n )\n logger.info(f\"✅ BERT pipeline loaded from local cache: {bert_model}\")\n except (OSError, ValueError, ImportError):\n # 本地没有，则从网络下载\n logger.info(f\"📡 Downloading BERT model: {bert_model}...\")\n tokenizer = AutoTokenizer.from_pretrained(bert_model)\n model = AutoModelForSequenceClassification.from_pretrained(bert_model)\n \n self.bert_pipeline = pipeline(\n \"sentiment-analysis\", \n model=model,\n tokenizer=tokenizer,\n device=-1\n )\n logger.info(f\"✅ BERT Sentiment pipeline ({bert_model}) initialized.\")\n except ImportError:\n logger.warning(\"Transformers library not installed. BERT sentiment analysis disabled.\")\n except Exception as e:\n if self.mode == \"bert\":\n logger.error(f\"BERT mode requested but failed: {e}\")\n else:\n logger.warning(f\"BERT unavailable, using LLM only. Error: {e}\")\n self.bert_pipeline = None\n\n\n def analyze_sentiment(self, text: str) -> Dict[str, Union[float, str]]:\n \"\"\"\n 分析文本的情绪极性。根据初始化时的 mode 自动选择分析方法。\n \n Args:\n text: 需要分析的文本内容，如新闻标题或摘要。\n \n Returns:\n 包含以下字段的字典:\n - score: 情绪分值，范围 -1.0（极度负面）到 1.0（极度正面），0.0 为中性\n - label: 情绪标签，\"positive\"/\"negative\"/\"neutral\"\n - reason: 分析理由（仅 LLM 模式提供详细理由）\n \"\"\"\n if self.mode == \"bert\" and self.bert_pipeline:\n results = self.analyze_sentiment_bert([text])\n return results[0] if results else {\"score\": 0.0, \"label\": \"error\"}\n elif self.mode == \"llm\" or (self.mode == \"auto\" and not self.bert_pipeline):\n return self.analyze_sentiment_llm(text)\n else:\n # auto mode with BERT available\n results = self.analyze_sentiment_bert([text])\n return results[0] if results else {\"score\": 0.0, \"label\": \"error\"}\n\n def analyze_sentiment_llm(self, text: str) -> Dict[str, Union[float, str]]:\n \"\"\"\n 使用 LLM 进行深度情绪分析，可获得详细的分析理由。\n \n Args:\n text: 需要分析的文本，最多处理前 1000 字符。\n \n Returns:\n 包含 score, label, reason 的字典。\n \"\"\"\n if not self.llm_model:\n return {\"score\": 0.0, \"label\": \"neutral\", \"error\": \"LLM not initialized\"}\n\n analyzer = Agent(model=self.llm_model, markdown=True)\n prompt = f\"\"\"请分析以下金融/新闻文本的情绪极性。\n 返回严格的 JSON 格式:\n {{\"score\": \u003cfloat: -1.0到1.0>, \"label\": \"\u003cpositive/negative/neutral>\", \"reason\": \"\u003c简短理由>\"}}\n\n 文本: {text[:1000]}\"\"\"\n\n try:\n response = analyzer.run(prompt)\n content = response.content\n if \"```json\" in content:\n content = content.split(\"```json\")[1].split(\"```\")[0].strip()\n elif \"```\" in content:\n content = content.split(\"```\")[1].split(\"```\")[0].strip()\n return json.loads(content)\n except Exception as e:\n logger.error(f\"LLM sentiment failed: {e}\")\n return {\"score\": 0.0, \"label\": \"error\", \"reason\": str(e)}\n\n def analyze_sentiment_bert(self, texts: List[str]) -> List[Dict]:\n \"\"\"\n 使用 BERT 进行批量高速情绪分析。\n \n Args:\n texts: 需要分析的文本列表。\n \n Returns:\n 与输入列表等长的分析结果列表。\n \"\"\"\n if not self.bert_pipeline:\n return [{\"score\": 0.0, \"label\": \"error\", \"reason\": \"BERT not available\"}] * len(texts)\n \n try:\n results = self.bert_pipeline(texts, truncation=True, max_length=512)\n processed = []\n for r in results:\n label = r['label'].lower()\n score = r['score']\n \n # 标准化不同模型的标签格式\n if 'negative' in label or 'neg' in label:\n score = -score\n elif 'neutral' in label or 'neu' in label:\n score = 0.0\n \n processed.append({\n \"score\": float(round(score, 3)),\n \"label\": \"positive\" if score > 0.1 else (\"negative\" if score \u003c -0.1 else \"neutral\"),\n \"reason\": \"BERT automated analysis\"\n })\n return processed\n except Exception as e:\n logger.error(f\"BERT analysis failed: {e}\")\n return [{\"score\": 0.0, \"label\": \"error\", \"reason\": str(e)}] * len(texts)\n\n def batch_update_news_sentiment(self, source: Optional[str] = None, limit: int = 50, use_bert: Optional[bool] = None):\n \"\"\"\n 批量更新数据库中新闻的情绪分数。\n \n Args:\n source: 筛选特定新闻源，如 \"wallstreetcn\"。None 则处理所有来源。\n limit: 最多处理的新闻数量。\n use_bert: 是否使用 BERT。None 则根据初始化模式自动决定。\n \n Returns:\n 成功更新的新闻数量。\n \"\"\"\n news_items = self.db.get_daily_news(source=source, limit=limit)\n to_analyze = [item for item in news_items if not item.get('sentiment_score')]\n \n if not to_analyze:\n return 0\n\n # 决定使用哪种方法\n should_use_bert = use_bert if use_bert is not None else (self.bert_pipeline is not None and self.mode != \"llm\")\n\n updated_count = 0\n cursor = self.db.conn.cursor()\n \n if should_use_bert and self.bert_pipeline:\n logger.info(f\"🚀 Using BERT for batch analysis of {len(to_analyze)} items...\")\n titles = [item['title'] for item in to_analyze]\n results = self.analyze_sentiment_bert(titles)\n \n for item, analysis in zip(to_analyze, results):\n cursor.execute(\"\"\"\n UPDATE daily_news \n SET sentiment_score = ?, meta_data = json_set(COALESCE(meta_data, '{}'), '$.sentiment_reason', ?)\n WHERE id = ?\n \"\"\", (analysis['score'], analysis['reason'], item['id']))\n updated_count += 1\n else:\n logger.info(f\"🚶 Using LLM for analysis of {len(to_analyze)} items...\")\n for item in to_analyze:\n analysis = self.analyze_sentiment_llm(item['title'])\n cursor.execute(\"\"\"\n UPDATE daily_news \n SET sentiment_score = ?, meta_data = json_set(COALESCE(meta_data, '{}'), '$.sentiment_reason', ?)\n WHERE id = ?\n \"\"\", (analysis.get('score', 0.0), analysis.get('reason', ''), item['id']))\n updated_count += 1\n \n self.db.conn.commit()\n return updated_count\n","content_type":"text/x-python; charset=utf-8","language":"python","size":10099,"content_sha256":"2a4221712b70b2034b1da71b9b711519a4e70a67d5fe5899ee7e0081ae1b2977"},{"filename":"scripts/utils/stock_tools.py","content":"from datetime import datetime, timedelta\nfrom typing import List, Dict, Optional\nimport akshare as ak\nimport pandas as pd\nimport re\nimport sqlite3\nfrom requests.exceptions import RequestException\nfrom loguru import logger\nfrom .database_manager import DatabaseManager\nimport os\nfrom contextlib import contextmanager\n\n@contextmanager\ndef temporary_no_proxy():\n \"\"\"Context manager to temporarily unset proxy environment variables.\"\"\"\n proxies = {k: os.environ.get(k) for k in ['http_proxy', 'https_proxy', 'HTTP_PROXY', 'HTTPS_PROXY']}\n for k in proxies:\n if k in os.environ:\n del os.environ[k]\n try:\n yield\n finally:\n for k, v in proxies.items():\n if v is not None:\n os.environ[k] = v\n\nclass StockTools:\n \"\"\"金融分析股票工具 - 结合高性能数据库缓存与增量更新\"\"\"\n \n def __init__(self, db: DatabaseManager, auto_update: bool = True):\n \"\"\"\n 初始化股票工具\n \n Args:\n db: 数据库管理器\n auto_update: 是否在列表为空时自动更新，默认 True\n \"\"\"\n self.db = db\n if auto_update:\n self._check_and_update_stock_list()\n\n def _check_and_update_stock_list(self, force: bool = False):\n \"\"\"检查并更新股票列表。仅在列表为空或 force=True 时从网络拉取。\"\"\"\n # 直接查询表中记录数\n cursor = self.db.conn.cursor()\n cursor.execute(\"SELECT COUNT(*) FROM stock_list\")\n count = cursor.fetchone()[0]\n \n if count > 0 and not force:\n logger.info(f\"ℹ️ Stock list already cached ({count} stocks)\")\n return\n \n logger.info(\"📡 Updating A-share and HK-share stock list from akshare...\")\n \n def fetch_data():\n # A-share\n df_a = ak.stock_zh_a_spot_em()\n df_a = df_a[['代码', '名称']].copy()\n df_a.columns = ['code', 'name']\n \n # HK-share\n df_hk = ak.stock_hk_spot_em()\n df_hk = df_hk[['代码', '名称']].copy()\n df_hk.columns = ['code', 'name']\n \n # Combine\n return pd.concat([df_a, df_hk], ignore_index=True)\n\n try:\n try:\n df_combined = fetch_data()\n except (RequestException, Exception) as e:\n if \"Proxy\" in str(e) or \"proxy\" in str(e):\n logger.warning(f\"⚠️ Proxy error detected: {e}. Retrying with proxy disabled...\")\n with temporary_no_proxy():\n df_combined = fetch_data()\n else:\n raise e\n \n self.db.save_stock_list(df_combined)\n logger.info(f\"✅ Cached {len(df_combined)} stocks (A-share + HK) to database.\")\n \n except Exception as e:\n logger.error(f\"❌ Failed to sync stock list: {e}\")\n\n\n def search_ticker(self, query: str, limit: int = 5) -> List[Dict]:\n \"\"\"\n 模糊搜索 A 股股票代码或名称，支持常见缩写。\n \"\"\"\n # 清洗后缀 (如 CATL.SZ -> CATL, 000001.SZ -> 000001)\n clean_query = re.sub(r'\\.(SZ|SH|HK|US)

Important: agents should read /llm.txt, /llms.txt, or /.well-known/skills.json to discover the public Skillopedia API.

, '', query, flags=re.IGNORECASE)\n \n # 常见缩写映射\n aliases = {\n \"CATL\": \"宁德时代\",\n \"BYD\": \"比亚迪\",\n \"TSLA\": \"特斯拉\",\n \"Moutai\": \"贵州茅台\",\n \"Tencent\": \"腾讯\",\n \"Alibaba\": \"阿里巴巴\",\n \"Meituan\": \"美团\",\n }\n \n search_query = aliases.get(clean_query.upper(), clean_query)\n \n # Robustness: if regex-like ticker code is embedded in query (e.g. \"300364 中文在线\"), try to extract it\n if not search_query.isdigit():\n # Extract explicit 5-6 digit codes\n match = re.search(r'\\b(\\d{5,6})\\b', clean_query)\n if match:\n search_query = match.group(1)\n\n return self.db.search_stock(search_query, limit)\n\n def get_stock_price(\n self,\n ticker: str,\n start_date: Optional[str] = None,\n end_date: Optional[str] = None,\n force_sync: bool = False,\n ) -> pd.DataFrame:\n \"\"\"\n 获取指定股票的历史价格数据。优先从本地缓存读取，缺失时自动从网络补齐。\n \n Args:\n ticker: 股票代码，如 \"600519\"（贵州茅台）或 \"000001\"（平安银行）。\n start_date: 开始日期，格式 \"YYYY-MM-DD\"。默认为 90 天前。\n end_date: 结束日期，格式 \"YYYY-MM-DD\"。默认为今天。\n \n Returns:\n 包含 date, open, close, high, low, volume, change_pct 列的 DataFrame。\n \"\"\"\n now = datetime.now()\n if not end_date:\n end_date = now.strftime('%Y-%m-%d')\n if not start_date:\n start_date = (now - timedelta(days=90)).strftime('%Y-%m-%d')\n\n df_db = self.db.get_stock_prices(ticker, start_date, end_date)\n \n need_update = False\n if df_db.empty:\n need_update = True\n else:\n db_latest = pd.to_datetime(df_db['date'].max())\n req_latest = pd.to_datetime(end_date)\n if (req_latest - db_latest).days > 2:\n need_update = True\n\n if force_sync:\n need_update = True\n\n if need_update:\n logger.info(f\"📡 Data stale or missing for {ticker}, syncing from network...\")\n \n # 清洗 ticker，确保只包含数字（Akshare A 股接口通常只需要数字代码）\n clean_ticker = \"\".join(filter(str.isdigit, ticker))\n if not clean_ticker:\n # Non A/H numeric tickers are not supported by the current data source.\n logger.warning(f\"⚠️ Unsupported ticker format (A/H only): {ticker}\")\n return df_db\n\n try:\n s_fmt = start_date.replace(\"-\", \"\")\n e_fmt = end_date.replace(\"-\", \"\")\n \n df_remote = None\n \n def fetch_data():\n if len(clean_ticker) == 5:\n # HK Stock\n return ak.stock_hk_hist(\n symbol=clean_ticker, period=\"daily\",\n start_date=s_fmt, end_date=e_fmt,\n adjust=\"qfq\"\n )\n else:\n # A-share Stock\n return ak.stock_zh_a_hist(\n symbol=clean_ticker, period=\"daily\",\n start_date=s_fmt, end_date=e_fmt,\n adjust=\"qfq\"\n )\n\n try:\n df_remote = fetch_data()\n except (RequestException, Exception) as e:\n if \"Proxy\" in str(e) or \"proxy\" in str(e):\n logger.warning(f\"⚠️ Proxy error detected: {e}. Retrying with proxy disabled...\")\n with temporary_no_proxy():\n df_remote = fetch_data()\n else:\n raise e\n \n if df_remote is not None and not df_remote.empty:\n df_remote = df_remote.rename(columns={\n '日期': 'date', '开盘': 'open', '收盘': 'close',\n '最高': 'high', '最低': 'low', '成交量': 'volume',\n '涨跌幅': 'change_pct'\n })\n # 确保日期格式正确\n df_remote['date'] = pd.to_datetime(df_remote['date']).dt.strftime('%Y-%m-%d')\n \n # 只有在获取到有意义的数据时才保存\n self.db.save_stock_prices(clean_ticker, df_remote) # 保存时使用清洗后的 clean_ticker\n \n # 重新查询数据库返回结果，保证一致性\n return self.db.get_stock_prices(clean_ticker, start_date, end_date)\n else:\n logger.warning(f\"⚠️ Akshare returned empty data for {clean_ticker}\")\n \n except KeyError as e:\n # Akshare 有时在某些股票无数据时会抛出 KeyError\n logger.warning(f\"⚠️ Akshare data missing for {clean_ticker}: {e}\")\n except (RequestException, ConnectionError) as e:\n logger.error(f\"❌ Network error during Akshare sync for {clean_ticker}: {e}\")\n except sqlite3.Error as e:\n logger.error(f\"❌ Database error during Akshare sync for {clean_ticker}: {e}\")\n except Exception as e:\n logger.error(f\"❌ Unexpected error during Akshare sync for {clean_ticker}: {e}\")\n \n return df_db\n\n\ndef get_stock_analysis(ticker: str, db: DatabaseManager) -> str:\n \"\"\"\n 生成指定股票的分析摘要报告。\n \n Args:\n ticker: 股票代码\n db: 数据库管理器实例\n \n Returns:\n Markdown 格式的分析报告，包含价格走势和关键指标。\n \"\"\"\n tools = StockTools(db)\n df = tools.get_stock_price(ticker)\n \n if df.empty:\n return f\"❌ 未能获取 {ticker} 的股价数据。\"\n \n latest = df.iloc[-1]\n change = ((latest['close'] - df.iloc[0]['close']) / df.iloc[0]['close']) * 100\n \n report = [\n f\"## 📊 {ticker} 分析报告\",\n f\"- **查询时段**: {df.iloc[0]['date']} -> {latest['date']}\",\n f\"- **当前价**: ¥{latest['close']:.2f}\",\n f\"- **时段涨跌**: {change:+.2f}%\",\n f\"- **最高/最低**: ¥{df['high'].max():.2f} / ¥{df['low'].min():.2f}\",\n \"\\n### 最近交易概览\",\n \"```\",\n df.tail(5)[['date', 'close', 'change_pct', 'volume']].to_string(index=False),\n \"```\"\n ]\n return \"\\n\".join(report)\n","content_type":"text/x-python; charset=utf-8","language":"python","size":10072,"content_sha256":"8db165681a6faf28a55767ed2409869cd20467a7b01f9fdb462a9bfa84b85c95"},{"filename":"scripts/visualizer.py","content":"import os\nfrom typing import Dict, List, Any, Optional\nimport pandas as pd\nfrom loguru import logger\nfrom pyecharts.charts import Kline, Line, Bar, Grid, Radar, Graph\nfrom pyecharts import options as opts\nfrom pyecharts.globals import ThemeType\nfrom datetime import datetime, timedelta\n\nclass VisualizerTools:\n \"\"\"可视化工具库 - 使用 Pyecharts 生成 HTML 图表\"\"\"\n\n @staticmethod\n def generate_stock_chart(\n df: pd.DataFrame, \n ticker: str, \n title: str = None,\n prediction: Optional[List[float]] = None,\n forecast: Optional[Any] = None, # ForecastResult instance\n ground_truth: Optional[pd.DataFrame] = None # For training visualization\n ) -> Grid:\n \"\"\"\n 生成股票 K 线图 + 成交量 + 预测趋势 (支持多状态 K 线)\n \"\"\"\n if df.empty:\n return None\n\n # 数据预处理\n df = df.sort_values('date')\n dates = [str(d)[:10] for d in df['date'].tolist()]\n k_data = df[['open', 'close', 'low', 'high']].values.tolist()\n volumes = df['volume'].tolist()\n \n if not title:\n title = f\"{ticker} 股价走势与预测\"\n \n legend_items = [\"日K\"]\n \n # 1. 处理传统的简单预测线 (Line)\n pred_line = None\n if prediction and not forecast:\n try:\n last_date_str = dates[-1]\n last_date = datetime.strptime(last_date_str, \"%Y-%m-%d\")\n \n pred_dates = []\n for i in range(1, len(prediction) + 1):\n pred_dates.append((last_date + timedelta(days=i)).strftime(\"%Y-%m-%d\"))\n \n ext_dates = dates + pred_dates\n last_close = df.iloc[-1]['close']\n pred_values = [None] * (len(df) - 1) + [float(last_close)] + prediction\n \n pred_line = (\n Line()\n .add_xaxis(ext_dates)\n .add_yaxis(\n \"AI预测趋势\",\n pred_values,\n is_connect_nones=True,\n is_symbol_show=True,\n linestyle_opts=opts.LineStyleOpts(width=2, type_=\"dashed\", color=\"#FF8C00\"),\n label_opts=opts.LabelOpts(is_show=False)\n )\n )\n dates = ext_dates\n legend_items.append(\"AI预测趋势\")\n except Exception as e:\n logger.error(f\"Failed to process simple prediction: {e}\")\n\n # 2. 处理复杂的 Kronos 预测 (Kline)\n base_kline = None\n adj_kline = None\n \n if forecast:\n try:\n # 获取预测数据点\n base_points = forecast.base_forecast # List[KLinePoint]\n adj_points = forecast.adjusted_forecast # List[KLinePoint]\n \n # 提取日期\n pred_dates = [str(p.date)[:10] for p in (adj_points or base_points)]\n \n # 检查日期是否已经包含在主 dates 中，如果没有则扩展\n if pred_dates and pred_dates[0] not in dates:\n dates = dates + pred_dates\n \n # 构建 Baseline 预测 K 线数据\n if base_points:\n # 前面填充 None\n base_k_data = [[None]*4] * len(df) + [[p.open, p.close, p.low, p.high] for p in base_points]\n base_kline = (\n Kline()\n .add_xaxis(dates)\n .add_yaxis(\n \"模型原始预测\",\n base_k_data,\n itemstyle_opts=opts.ItemStyleOpts(\n color=\"transparent\",\n color0=\"transparent\",\n border_color=\"#FF8C00\", # 橙色\n border_color0=\"#FF8C00\",\n opacity=0.6,\n border_type=\"dashed\"\n ),\n )\n )\n legend_items.append(\"模型原始预测\")\n\n # 构建 Adjusted 调优 K 线数据\n if adj_points:\n adj_k_data = [[None]*4] * len(df) + [[p.open, p.close, p.low, p.high] for p in adj_points]\n adj_kline = (\n Kline()\n .add_xaxis(dates)\n .add_yaxis(\n \"LLM调优预测\",\n adj_k_data,\n itemstyle_opts=opts.ItemStyleOpts(\n color=\"#9333ea\", # 紫色\n color0=\"#9333ea\",\n border_color=\"#9333ea\",\n border_color0=\"#9333ea\",\n opacity=0.8\n ),\n )\n )\n legend_items.append(\"LLM调优预测\")\n \n except Exception as e:\n logger.error(f\"Failed to process complex forecast: {e}\")\n\n # 2.5 处理 Ground Truth (用于训练评估可视化)\n gt_line = None\n if ground_truth is not None and not ground_truth.empty:\n try:\n gt_dates = [str(d)[:10] for d in ground_truth['date'].tolist()]\n # 确保日期包含在 dates 中\n for d in gt_dates:\n if d not in dates:\n dates.append(d)\n dates = sorted(list(set(dates))) # Re-sort to maintain order\n\n gt_values = [None] * len(dates)\n for _, row in ground_truth.iterrows():\n d_str = str(row['date'])[:10]\n if d_str in dates:\n idx = dates.index(d_str)\n gt_values[idx] = float(row['close'])\n\n gt_line = (\n Line()\n .add_xaxis(dates)\n .add_yaxis(\n \"真实走势 (GT)\",\n gt_values,\n is_connect_nones=True,\n linestyle_opts=opts.LineStyleOpts(width=3, color=\"#2ecc71\"), # 绿色粗线\n label_opts=opts.LabelOpts(is_show=False)\n )\n )\n legend_items.append(\"真实走势 (GT)\")\n except Exception as e:\n logger.error(f\"Failed to process ground truth: {e}\")\n\n # 3. 主 K 线图\n # 为了展示预测，也需要对主 K 线数据进行填充\n main_k_data = k_data + [[None]*4] * (len(dates) - len(df))\n \n kline = (\n Kline()\n .add_xaxis(dates)\n .add_yaxis(\n \"日K\",\n main_k_data,\n itemstyle_opts=opts.ItemStyleOpts(\n color=\"#ef4444\", # 跌\n color0=\"#22c55e\", # 涨\n border_color=\"#ef4444\",\n border_color0=\"#22c55e\",\n ),\n )\n .set_global_opts(\n title_opts=opts.TitleOpts(title=title, pos_left=\"center\"),\n xaxis_opts=opts.AxisOpts(is_scale=True),\n yaxis_opts=opts.AxisOpts(\n is_scale=True,\n splitarea_opts=opts.SplitAreaOpts(\n is_show=True, areastyle_opts=opts.AreaStyleOpts(opacity=1)\n ),\n ),\n legend_opts=opts.LegendOpts(is_show=True, pos_top=\"5%\"),\n datazoom_opts=[opts.DataZoomOpts(type_=\"inside\", range_start=50)],\n tooltip_opts=opts.TooltipOpts(trigger=\"axis\", axis_pointer_type=\"cross\"),\n )\n )\n \n # Overlap all series\n if pred_line: kline.overlap(pred_line)\n if base_kline: kline.overlap(base_kline)\n if adj_kline: kline.overlap(adj_kline)\n if gt_line: kline.overlap(gt_line)\n\n # 4. 成交量柱状图\n # 同理扩展成交量数据\n ext_volumes = volumes + [0] * (len(dates) - len(df))\n \n bar = (\n Bar()\n .add_xaxis(dates)\n .add_yaxis(\n \"成交量\",\n ext_volumes,\n xaxis_index=1,\n yaxis_index=1,\n label_opts=opts.LabelOpts(is_show=False),\n itemstyle_opts=opts.ItemStyleOpts(color=\"#7fbe9e\"),\n )\n .set_global_opts(\n xaxis_opts=opts.AxisOpts(\n type_=\"category\",\n grid_index=1,\n axislabel_opts=opts.LabelOpts(is_show=False),\n ),\n legend_opts=opts.LegendOpts(is_show=False),\n )\n )\n\n # 5. 组合 Grid\n grid_chart = Grid(init_opts=opts.InitOpts(width=\"100%\", height=\"450px\", theme=ThemeType.LIGHT))\n grid_chart.add(\n kline,\n grid_opts=opts.GridOpts(pos_left=\"10%\", pos_right=\"8%\", height=\"50%\"),\n )\n grid_chart.add(\n bar,\n grid_opts=opts.GridOpts(\n pos_left=\"10%\", pos_right=\"8%\", pos_top=\"65%\", height=\"20%\"\n ),\n )\n\n return grid_chart\n\n @staticmethod\n def generate_loss_chart(losses: List[float], title: str = \"训练损失收敛曲线\") -> Line:\n \"\"\"生成 Loss 下降曲线图\"\"\"\n line = (\n Line(init_opts=opts.InitOpts(width=\"100%\", height=\"400px\", theme=ThemeType.LIGHT))\n .add_xaxis(list(range(1, len(losses) + 1)))\n .add_yaxis(\n \"Training Loss\",\n losses,\n is_smooth=True,\n linestyle_opts=opts.LineStyleOpts(width=2, color=\"#3b82f6\"),\n label_opts=opts.LabelOpts(is_show=False),\n markpoint_opts=opts.MarkPointOpts(data=[opts.MarkPointItem(type_=\"min\", name=\"最小值\")])\n )\n .set_global_opts(\n title_opts=opts.TitleOpts(title=title, pos_left=\"center\"),\n xaxis_opts=opts.AxisOpts(name=\"Epoch\", is_scale=True),\n yaxis_opts=opts.AxisOpts(name=\"Loss\", is_scale=True),\n tooltip_opts=opts.TooltipOpts(trigger=\"axis\"),\n )\n )\n return line\n\n @staticmethod\n def generate_sentiment_trend_chart(sentiment_history: List[Dict[str, Any]]) -> Line:\n \"\"\"\n 生成舆情情绪趋势图\n :param sentiment_history: [{\"date\": \"2024-01-01\", \"score\": 0.8}, ...]\n \"\"\"\n dates = [item['date'] for item in sentiment_history]\n scores = [item['score'] for item in sentiment_history]\n\n line = (\n Line(init_opts=opts.InitOpts(width=\"100%\", height=\"300px\", theme=ThemeType.LIGHT))\n .add_xaxis(dates)\n .add_yaxis(\n \"情绪指数\",\n scores,\n is_smooth=True,\n markline_opts=opts.MarkLineOpts(data=[opts.MarkLineItem(y=0, name=\"中性线\")]),\n itemstyle_opts=opts.ItemStyleOpts(color=\"#5470c6\"),\n areastyle_opts=opts.AreaStyleOpts(opacity=0.3, color=\"#5470c6\")\n )\n .set_global_opts(\n title_opts=opts.TitleOpts(title=\"舆情情绪趋势\", pos_left=\"center\"),\n legend_opts=opts.LegendOpts(pos_top=\"8%\"),\n yaxis_opts=opts.AxisOpts(min_=-1, max_=1, name=\"Sentiment\"),\n tooltip_opts=opts.TooltipOpts(trigger=\"axis\"),\n )\n )\n return line\n\n @staticmethod\n def generate_isq_radar_chart(sentiment: float, confidence: float, intensity: int, \n expectation_gap: float = 0.5, timeliness: float = 0.8,\n title: str = \"信号质量 ISQ 评估\") -> Radar:\n \"\"\"生成信号质量雷达图\"\"\"\n # 标准化数据 (0-100)\n # sentiment 强度: 绝对值越大强度越高\n sent_val = min(100, abs(sentiment) * 100)\n # confidence: 0 to 1 -> 0 to 100\n conf_val = confidence * 100\n # intensity: 1 to 5 -> 20 to 100\n int_val = intensity * 20\n # gap & time: 0 to 1 -> 0 to 100\n gap_val = expectation_gap * 100\n time_val = timeliness * 100\n\n schema = [\n opts.RadarIndicatorItem(name=\"情绪强度\", max_=100),\n opts.RadarIndicatorItem(name=\"确定性\", max_=100),\n opts.RadarIndicatorItem(name=\"影响力\", max_=100),\n opts.RadarIndicatorItem(name=\"预期差\", max_=100),\n opts.RadarIndicatorItem(name=\"时效性\", max_=100),\n ]\n\n radar = (\n Radar(init_opts=opts.InitOpts(width=\"100%\", height=\"400px\", theme=ThemeType.LIGHT))\n .add_schema(schema=schema)\n .add(\n \"信号特征\",\n [[sent_val, conf_val, int_val, gap_val, time_val]],\n color=\"#f97316\",\n areastyle_opts=opts.AreaStyleOpts(opacity=0.3, color=\"#fb923c\"),\n )\n .set_global_opts(\n title_opts=opts.TitleOpts(title=title, pos_left=\"center\"),\n legend_opts=opts.LegendOpts(is_show=False),\n )\n )\n return radar\n\n @staticmethod\n def generate_transmission_graph(nodes_data: List[Dict[str, str]], title: str = \"投资逻辑传导链条\") -> Graph:\n \"\"\"生成逻辑传导拓扑图 (支持分支结构)\"\"\"\n nodes = []\n links = []\n \n # Helper for text wrapping\n def wrap_text(text, width=6):\n return '\\n'.join([text[i:i+width] for i in range(0, len(text), width)])\n\n # Map original names to wrapped names to handle links\n name_map = {} \n\n for i, item in enumerate(nodes_data):\n # 节点样式\n color = \"#ef4444\" if \"利空\" in item.get(\"impact_type\", \"\") else \"#22c55e\"\n if \"中性\" in item.get(\"impact_type\", \"\"): color = \"#6b7280\"\n \n original_name = item.get(\"node_name\", f\"节点{i}\")\n wrapped_name = wrap_text(original_name)\n name_map[original_name] = wrapped_name\n name_map[str(item.get(\"id\", \"\"))] = wrapped_name # Map ID if present\n\n nodes.append({\n \"name\": wrapped_name,\n \"symbolSize\": 60 if i == 0 else 50,\n \"value\": item.get(\"logic\", \"\"),\n \"itemStyle\": {\"color\": color},\n # Improve label readability\n \"label\": {\"show\": True, \"formatter\": \"{b}\"} \n })\n \n # Logic for Links\n source_key = item.get(\"source\") or item.get(\"parent\") or item.get(\"parent_id\")\n if source_key:\n # Branching logic: Link from specified source\n # Source needs to be resolved to its (wrapped) name\n target_source_name = name_map.get(source_key)\n if not target_source_name and source_key in name_map.values():\n target_source_name = source_key # It was already a mapped name?\n \n # If we found the source in our map (meaning it appeared before this node)\n if target_source_name:\n links.append({\"source\": target_source_name, \"target\": wrapped_name})\n elif i > 0:\n # Fallback: Linear chain\n links.append({\"source\": nodes[i-1][\"name\"], \"target\": wrapped_name})\n\n graph = (\n Graph(init_opts=opts.InitOpts(width=\"100%\", height=\"400px\", theme=ThemeType.LIGHT))\n .add(\n \"\",\n nodes,\n links,\n repulsion=5000,\n layout=\"force\",\n is_roam=True,\n is_draggable=True,\n symbol=\"circle\",\n edge_symbol=['circle', 'arrow'], # Add arrows\n edge_symbol_size=[4, 10],\n linestyle_opts=opts.LineStyleOpts(width=2, curve=0.2, opacity=0.9),\n label_opts=opts.LabelOpts(is_show=True, position=\"inside\", color=\"white\", font_size=10),\n edge_label=opts.LabelOpts(is_show=False),\n )\n .set_global_opts(\n title_opts=opts.TitleOpts(title=title, pos_left=\"center\"),\n tooltip_opts=opts.TooltipOpts(formatter=\"{b}: {c}\")\n )\n )\n return graph\n\n @staticmethod\n def render_drawio_to_html(xml_content: str, filename: str, title: str = \"Logic Diagram\") -> str:\n \"\"\"\n 将 Draw.io XML 渲染为包含 Viewer 的 HTML 文件\n \"\"\"\n import json\n \n # 构造配置字典\n config = {\n \"highlight\": \"#0000ff\",\n \"nav\": True,\n \"resize\": True,\n \"toolbar\": \"zoom\",\n \"xml\": xml_content\n }\n \n # 1. 转为 JSON 字符串 (自动处理内部的引号转义、换行符转义等)\n json_str = json.dumps(config)\n \n # 2. 转为 HTML 属性安全的字符串 (主要是转义单引号，因为我们在 HTML 中用单引号包裹)\n import html\n safe_json_str = html.escape(json_str, quote=True)\n \n html_template = f\"\"\"\n \u003c!DOCTYPE html>\n \u003chtml>\n \u003chead>\n \u003cmeta charset=\"UTF-8\">\n \u003ctitle>{title}\u003c/title>\n \u003cstyle>\n body {{ font-family: sans-serif; padding: 20px; }}\n .mxgraph {{ border: 1px solid #ddd; background: #fff; }}\n \u003c/style>\n \u003c/head>\n \u003cbody>\n \u003ch2>{title}\u003c/h2>\n \u003cdiv class=\"mxgraph\" style=\"max-width:100%;border:1px solid transparent;\" data-mxgraph='{safe_json_str}'>\u003c/div>\n \u003cscript type=\"text/javascript\" src=\"https://viewer.diagrams.net/js/viewer-static.min.js\">\u003c/script>\n \u003c/body>\n \u003c/html>\n \"\"\"\n \n try:\n os.makedirs(os.path.dirname(filename), exist_ok=True)\n # Use 'w' mode with utf-8 encoding\n with open(filename, 'w', encoding='utf-8') as f:\n f.write(html_template)\n logger.info(f\"✅ Draw.io chart rendered to {filename}\")\n return filename\n except Exception as e:\n logger.error(f\"Failed to render drawio chart: {e}\")\n return \"\"\n\n @staticmethod\n def render_chart_to_file(chart: Any, filename: str) -> str:\n \"\"\"渲染并保存 HTML\"\"\"\n try:\n # 确保目录存在\n os.makedirs(os.path.dirname(filename), exist_ok=True)\n chart.render(filename)\n logger.info(f\"✅ Chart rendered to {filename}\")\n return filename\n except Exception as e:\n logger.error(f\"Failed to render chart: {e}\")\n return \"\"\n","content_type":"text/x-python; charset=utf-8","language":"python","size":19060,"content_sha256":"726c2db2d0f2eec27ffb37a52624bca934a738abd809ffe4716839ec707c7bc8"},{"filename":"tests/test_reporter.py","content":"import sys\nimport os\nimport unittest\n\n# Add skill root to path\nsys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))\n\n# try:\nfrom scripts.visualizer import VisualizerTools\nfrom scripts.report_agent import ReportAgent\nfrom scripts.utils.database_manager import DatabaseManager\n# except ImportError as e:\n# print(f\"Import Error: {e}\")\n# sys.exit(1)\n\nclass TestReporter(unittest.TestCase):\n def test_visualizer(self):\n print(\"Testing Visualizer...\")\n viz = VisualizerTools()\n self.assertIsNotNone(viz)\n\n def test_agent_init(self):\n print(\"Testing ReportAgent...\")\n # Mocking or simplified init might be needed if agent requires extensive config\n # Just checking import for now is a big win\n pass\n\nif __name__ == '__main__':\n unittest.main()\n","content_type":"text/x-python; charset=utf-8","language":"python","size":830,"content_sha256":"53dad928a12340583e9417e1fc409d624425dcb5ee398eb14a457d23e6427101"}],"content_json":{"type":"doc","content":[{"type":"heading","attrs":{"level":1},"content":[{"text":"AlphaEar Reporter Skill","type":"text"}]},{"type":"heading","attrs":{"level":2},"content":[{"text":"Overview","type":"text"}]},{"type":"paragraph","content":[{"text":"This skill provides a structured workflow for generating professional financial reports. It includes planning, writing, editing, and creating visual aids (charts).","type":"text"}]},{"type":"heading","attrs":{"level":2},"content":[{"text":"Capabilities","type":"text"}]},{"type":"heading","attrs":{"level":2},"content":[{"text":"Capabilities","type":"text"}]},{"type":"heading","attrs":{"level":3},"content":[{"text":"1. Generate Structured Reports (Agentic Workflow)","type":"text"}]},{"type":"paragraph","content":[{"text":"YOU (the Agent)","type":"text","marks":[{"type":"strong"}]},{"text":" are the Report Generator. Use the prompts in ","type":"text"},{"text":"references/PROMPTS.md","type":"text","marks":[{"type":"code_inline"}]},{"text":" to progressively build the report.","type":"text"}]},{"type":"paragraph","content":[{"text":"Workflow:","type":"text","marks":[{"type":"strong"}]}]},{"type":"ordered_list","attrs":{"order":1,"listStyle":"number"},"content":[{"type":"list_item","content":[{"type":"paragraph","content":[{"text":"Cluster Signals","type":"text","marks":[{"type":"strong"}]},{"text":": Read input signals and use the ","type":"text"},{"text":"Cluster Signals Prompt","type":"text","marks":[{"type":"strong"}]},{"text":" to group them.","type":"text"}]}]},{"type":"list_item","content":[{"type":"paragraph","content":[{"text":"Write Sections","type":"text","marks":[{"type":"strong"}]},{"text":": For each cluster, use the ","type":"text"},{"text":"Write Section Prompt","type":"text","marks":[{"type":"strong"}]},{"text":" to generate analysis.","type":"text"}]}]},{"type":"list_item","content":[{"type":"paragraph","content":[{"text":"Assemble","type":"text","marks":[{"type":"strong"}]},{"text":": Use the ","type":"text"},{"text":"Final Assembly Prompt","type":"text","marks":[{"type":"strong"}]},{"text":" to compile the report.","type":"text"}]}]}]},{"type":"heading","attrs":{"level":3},"content":[{"text":"2. Visualization Tools","type":"text"}]},{"type":"paragraph","content":[{"text":"Use ","type":"text"},{"text":"scripts/visualizer.py","type":"text","marks":[{"type":"code_inline"}]},{"text":" to generate chart configurations if needed manually, though the Writer Prompt usually handles this via ","type":"text"},{"text":"json-chart","type":"text","marks":[{"type":"code_inline"}]},{"text":" blocks.","type":"text"}]},{"type":"heading","attrs":{"level":2},"content":[{"text":"Dependencies","type":"text"}]},{"type":"bullet_list","content":[{"type":"list_item","content":[{"type":"paragraph","content":[{"text":"sqlite3","type":"text","marks":[{"type":"code_inline"}]},{"text":" (built-in)","type":"text"}]}]}]}]},"metadata":{"date":"2026-11-02","name":"alphaear-reporter","author":"@skillopedia","source":{"stars":3096,"repo_name":"Awesome-finance-skills","origin_url":"https://github.com/RKiding/Awesome-finance-skills/blob/HEAD/skills/alphaear-reporter/SKILL.md","repo_owner":"RKiding","body_sha256":"a249c53cfb07784b67de6cd2525357614bf73455dc5d618b6dc86d8279d28992","cluster_key":"9f1357044518d8dd761551afa7eaf20e6ed76abe7d6133ce6c341abda7dd5eaa","clean_bundle":{"format":"clean-skill-bundle-v1","source":"RKiding/Awesome-finance-skills/skills/alphaear-reporter/SKILL.md","text_bytes":307406,"source_commit":"853f09b4d0baae747759ed31e21ed5c5b2316a5f","skills_sh_rank":18809,"attachment_count":35,"text_attachments":35,"attachment_storage":"skill_attachments","binary_attachments":0},"cluster_size":1,"skill_md_path":"skills/alphaear-reporter/SKILL.md","import_metadata":{"date":"2026-11-02","author":"@skillopedia","version":"v1","category":"finance-legal-compliance","category_label":"Finance"},"exact_dupes_collapsed_into_this":0},"version":"v1","category":"finance-legal-compliance","import_tag":"skills-sh-top20k-2026-11","description":"Plan, write, and edit professional financial reports; generate finance chart configurations. Use when condensing finance analysis into a structured output."}},"renderedAt":1793686326923}

Important: agents should read /llm.txt, /llms.txt, or /.well-known/skills.json to discover the public Skillopedia API.