Upload folder using huggingface_hub

modules/title_generator/__init__.py

@@ -0,0 +1,3 @@
+"""
+Title Generator package initialization.
+"""

modules/title_generator/create_index.py

@@ -0,0 +1,65 @@
+import argparse
+import os
+import sys
+import json
+import logging
+from tqdm import tqdm
+
+# 添加项目根目录到Python路径
+root_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+sys.path.append(root_dir)
+
+# 然后使用相对于项目根目录的导入
+from modules.title_generator.title_generator import RagTitleGenerator
+
+# 配置日志
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger("BuildTitleIndex")
+
+def process(data,
+            index_path: str='faiss_infographics.index',
+            data_path: str='infographics_data.npy',
+            embed_model_path: str='',
+            force: bool=False):
+    try:
+        # 检查索引文件是否存在
+        if os.path.exists(index_path) and not force:
+            logger.info(f"索引文件 {index_path} 已存在，跳过创建。使用 --force 参数强制重建。")
+            return True
+            
+        logger.info("Initializing RagTitleGenerator...")
+        generator = RagTitleGenerator(
+            index_path=index_path,
+            data_path=data_path,
+            embed_model_path=embed_model_path
+        )
+        
+        logger.info("Building FAISS index...")
+        generator.build_faiss_index(data)
+        logger.info("Index built and saved successfully.")
+        return True
+    except Exception as e:
+        logger.error(f"构建索引失败: {str(e)}")
+        return False
+
+def main(force: bool=False):
+    parser = argparse.ArgumentParser(description="Build FAISS index for title generation")
+    parser.add_argument('--data', type=str, required=True, help='Path to training data JSON file')
+    parser.add_argument('--index_path', type=str, default='faiss_infographics.index', help='Path to store FAISS index')
+    parser.add_argument('--data_path', type=str, default='infographics_data.npy', help='Path to store embedding + title data')
+    parser.add_argument('--embed_model_path', type=str, default='', help='Path to sentence embedding model (optional)')
+    parser.add_argument('--force', action='store_true', help='Force rebuild even if index exists')
+
+    args = parser.parse_args()
+
+    process(args.data, args.index_path, args.data_path, args.embed_model_path, force or args.force)
+
+if __name__ == "__main__":
+    try:
+        main()
+    except Exception as e:
+        logger.error(f"程序执行失败: {str(e)}")
+        sys.exit(1)

modules/title_generator/title_generator.py

@@ -0,0 +1,459 @@
+import os
+import json
+import faiss
+import numpy as np
+from tqdm import tqdm
+import argparse
+import logging
+from typing import Any, Dict, List, Tuple, Union
+from openai import OpenAI
+from utils.model_loader import ModelLoader
+
+# 配置日志
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger("RagTitleGenerator")
+
+_GENERATOR_CACHE = {}
+
+
+def _has_required_titles(data: Dict[str, Any]) -> bool:
+    titles = data.get("titles")
+    return (
+        isinstance(titles, dict)
+        and bool(str(titles.get("main_title", "")).strip())
+        and bool(str(titles.get("sub_title", "")).strip())
+    )
+
+
+def _get_generator(
+    index_path: str,
+    data_path: str,
+    embed_model_path: str,
+    api_key: str,
+    base_url: str
+) -> "RagTitleGenerator":
+    key = (index_path, data_path, embed_model_path, api_key, base_url)
+    generator = _GENERATOR_CACHE.get(key)
+    if generator is None:
+        generator = RagTitleGenerator(
+            index_path=index_path,
+            data_path=data_path,
+            embed_model_path=embed_model_path,
+            api_key=api_key,
+            base_url=base_url
+        )
+        _GENERATOR_CACHE[key] = generator
+    return generator
+
+
+class RagTitleGenerator:
+    """
+    RagTitleGenerator is a class that handles both:
+    1. Building/maintaining the FAISS index from a given dataset.
+    2. Generating infographic titles and descriptions for a single data item using RAG.
+    """
+
+    def __init__(
+        self,
+        index_path: str = "faiss_infographics.index",
+        data_path: str = "infographics_data.npy",
+        embed_model_path="",
+        api_key: str="",
+        base_url: str=""
+    ) -> None:
+        """
+        Initialize the RagTitleGenerator. It attempts to load any existing FAISS index
+        and corresponding training data from disk. If no index file is found, the index
+        will be None until build_faiss_index is called.
+
+        Args:
+            index_path (str, optional): Path where the FAISS index file is or will be stored.
+            data_path (str, optional): Path where the training data embeddings are stored.
+            embed_model_path: Custom embedding model. If None, defaults to SentenceTransformer("all-MiniLM-L6-v2").
+        """
+        self.index_path = index_path
+        self.data_path = data_path
+        if embed_model_path:
+            self.embed_model = ModelLoader.get_model(embed_model_path)
+        else:
+            print("fuck")
+        #else:
+        #    self.embed_model = SentenceTransformer("all-MiniLM-L6-v2")
+
+        self.index = None
+        self.training_data = []  # List of tuples: (input_text, title, description)
+
+        print(api_key, base_url)
+        self.client = OpenAI(
+            api_key=api_key,
+            base_url=base_url
+            # "https://aihubmix.com/v1"
+        )
+
+        # Try loading an existing FAISS index and data
+        if os.path.exists(self.index_path) and os.path.exists(self.data_path):
+            print("Load existing FAISS index from disk.")
+            self.index = faiss.read_index(self.index_path)
+            with open(self.data_path, "rb") as f:
+                self.training_data = np.load(f, allow_pickle=True).tolist()
+        else:
+            print("No existing FAISS index found; you can build a new one via build_faiss_index().")
+
+    def build_faiss_index(self, data_json_path):
+        """Build a FAISS index from a JSON dataset of chart data."""
+        try:
+            if not os.path.exists(data_json_path):
+                raise ValueError(f"Provided data_json_path {data_json_path} does not exist.")
+
+            # Remove existing index
+            self.clear_faiss_data()
+
+            # Load the entire dataset from JSON
+            with open(data_json_path, "r", encoding="utf-8") as f:
+                dataset = json.load(f)
+            print("dataset", type(dataset))
+            # Build the FAISS index from scratch
+            if isinstance(dataset, dict):
+                for name, details in tqdm(dataset.items(), desc="Building new FAISS index"):
+                    processed_text = self.process_single_data(details)
+                    title = details.get("metadata", {}).get("title", "")
+                    description = details.get("metadata", {}).get("description", "")
+                    self.add_training_data(processed_text, title, description)
+            else:
+                raise ValueError("Dataset must be a dictionary")
+
+        except Exception as e:
+            logger.error(f"处理记录时出错: {str(e)}")
+            raise
+
+    def clear_faiss_data(self) -> None:
+        """Remove existing FAISS index and training data from disk."""
+        if os.path.exists(self.index_path):
+            os.remove(self.index_path)
+        if os.path.exists(self.data_path):
+            os.remove(self.data_path)
+        self.index = None
+        self.training_data = []
+        print("Original FAISS has been cleared.")
+
+    def process_single_data(
+        self,
+        data: Dict[str, Any]
+    ) -> str:
+        """
+        Convert a single data dictionary into a textual representation for retrieval and generation.
+        This includes metadata, chart_type, datafacts, etc.
+
+        Args:
+            data (Dict[str, Any]): A dictionary containing chart data, metadata, etc.
+
+        Returns:
+            str: A concatenated text representation of the data.
+        """
+        metadata = data.get("metadata", {})
+        chart_type = data.get("chart_type", [])
+        datafacts = data.get("datafacts", [])
+        data_columns = data["data"].get("columns", [])
+        chart_data = data["data"].get("data", [])
+
+        main_insight = metadata.get("main_insight", "")
+
+        # Convert chart type to text
+        chart_type_text = "Chart Type: " + ", ".join(chart_type) + "\n" if chart_type else ""
+
+        # Convert datafacts to text
+        datafacts_text = "Data Facts:\n"
+        if datafacts:
+            for fact in datafacts:
+                annotation = fact.get("annotation", "")
+                if annotation:
+                    datafacts_text += f"- {annotation}\n"
+        else:
+            datafacts_text = ""
+
+        # Main insight
+        main_insight_text = f"Main Insight: {main_insight}\n" if main_insight else ""
+
+        # Data columns
+        column_text = "Columns: "
+        if data_columns:
+            column_names = [col.get("name", "") for col in data_columns]
+            # column_names = [f"{col.get('name', '')} ({col.get('description', '')})" for col in data_columns]
+            column_text += ", ".join(column_names) + "\n"
+        else:
+            column_text = ""
+
+        # Data sample (limit to first few rows)
+        data_text = "Data Sample:\n"
+        for row in chart_data:
+            row_text = ", ".join([f"{k}: {v}" for k, v in row.items()])
+            data_text += f"{row_text}\n"
+
+        final_text = (
+            f"{chart_type_text}\n"
+            f"{datafacts_text}\n"
+            f"{main_insight_text}\n"
+            f"{column_text}\n"
+            f"{data_text}"
+        )
+        return final_text
+
+    def add_training_data(
+        self,
+        input_text: str,
+        title: str,
+        description: str
+    ) -> None:
+        """
+        Add a single training sample into the in-memory list and FAISS index.
+        After adding, it writes the updated index and data to disk.
+
+        Args:
+            input_text (str): Concatenated text derived from chart data.
+            title (str): Ground truth or known title from the data's metadata.
+            description (str): Ground truth or known description from the data's metadata.
+        """
+        self.training_data.append((input_text, title, description))
+
+        embedding = self.embed_model.encode([input_text])
+
+        if self.index is None:
+            dim = embedding.shape[1]
+            self.index = faiss.IndexFlatL2(dim)
+
+        self.index.add(np.array(embedding))
+
+        faiss.write_index(self.index, self.index_path)
+
+        with open(self.data_path, "wb") as f:
+            np.save(f, np.array(self.training_data, dtype=object))
+
+    def retrieve_similar(
+        self,
+        new_input: str,
+        topk: int = 7
+    ) -> List[Tuple[str, str, str]]:
+        """
+        Retrieve the top-k most similar training samples from the FAISS index.
+
+        Args:
+            new_input (str): The new text query to encode.
+            topk (int, optional): Number of similar samples to retrieve.
+
+        Returns:
+            List[Tuple[str, str, str]]: A list of (input_text, title, description).
+        """
+        if self.index is None or len(self.training_data) == 0:
+            return []
+
+        new_embedding = self.embed_model.encode([new_input])
+        topk = min(topk, len(self.training_data))
+        _, I = self.index.search(np.array(new_embedding), k=topk)
+
+        retrieved_list = []
+        for idx in I[0]:
+            data = self.training_data[idx]
+            retrieved_list.append((data[0], data[1], data[2]))
+
+        return retrieved_list
+
+    def generate_title_description(
+        self,
+        data: Dict[str, Any],
+        topk: int = 7
+    ) -> Tuple[str, str]:
+        """
+        Generate a title and subtitle for a single data dictionary using RAG.
+
+        Args:
+            data (Dict[str, Any]): The data dictionary containing metadata, chart_type, datafacts, etc.
+            topk (int, optional): Number of similar examples to retrieve for prompt augmentation.
+
+        Returns:
+            Tuple[str, str]: (generated_title, generated_description)
+        """
+        processed_text = self.process_single_data(data)
+
+        retrieved_examples = self.retrieve_similar(processed_text, topk=topk) if topk > 0 else []
+
+        max_title_words = 8
+        max_description_words = 13
+        for _, rt, rd in retrieved_examples:
+            rt = rt or ""
+            rd = rd or ""
+            max_title_words = max(max_title_words, len(rt.split()))
+            max_description_words = max(max_description_words, len(rd.split()))
+
+        example_text = ""
+        if retrieved_examples:
+            example_text += "Here are some similar examples:\n"
+            for i, (r_data, rt, rd) in enumerate(retrieved_examples, start=1):
+                example_text += (
+                    f"\n[Similar Example {i}]\n"
+                    f"Input Data:\n{r_data}\n"
+                    f"Title: {rt if rt else ''}\n"
+                    f"Description: {rd if rd else ''}\n"
+                )
+        title_prompt = (
+            f"{example_text}\n"
+            "Based on the above examples (if any), please generate a clear and concise TITLE for the following data.\n"
+            f"{processed_text}\n\n"
+            "Important instructions:\n"
+            "1. The title should focus on the most significant feature of the data. "
+            "You can choose one or more key insights from the Data Facts that best "
+            "illustrate the issue, or identify the most notable feature yourself.\n"
+            # "2. The title should focus solely on what the data is about, without analyzing specific "
+            # "data characteristics, trends, distributions, or comparisons.\n"
+            f"2. The title should be strictly under {max_title_words} words.\n"
+            "3. Use exact terminology from the data sources.\n"
+            "4. Do NOT use these verbs: show, reveal, illustrate, analyze.\n"
+            "5. ONLY return the title as a string, no extra text."
+        )
+
+        response_title = self.client.chat.completions.create(
+            model="gpt-4o-mini",
+            messages=[
+                {"role": "system", "content": "You are an AI assistant that generates infographic titles."},
+                {"role": "user", "content": title_prompt}
+            ]
+        )
+        generated_title = response_title.choices[0].message.content
+        generated_title = generated_title.strip() if generated_title else ""
+
+        description_prompt = (
+            f"{example_text}\n"
+            "Based on the above examples (if any), please generate a precise DESCRIPTION for the following data.\n"
+            f"{processed_text}\n\n"
+            "Important instructions:\n"
+            "1. Do NOT describe statistical properties (e.g., highest/lowest values, changes over time, "
+            "ratios, percentages). Simply summarize what the dataset reports.\n"
+            "2. Use one of the following structured templates where applicable (choose the highest-priority one that fits):\n"
+            "   - Share/Percentage of [group] (who [action/characteristic]) (by [region/timeframe] (in [units])).\n"
+            "   - Number/Total/Amount of [entity] (in [region/timeframe]) (, measured in [units]).\n"
+            "   - Top/Leading N [entities] by [indicator] (, in [timeframe/region]).\n"
+            "   - [Indicator] for [group/topic] (in [region/timeframe]) (, measured in [units]).\n"
+            f"3. The description should be strictly under {max_description_words} words.\n"
+            "4. Use exact terminology from the data sources.\n"
+            "5. Do NOT use these verbs: show, reveal, illustrate, analyze.\n"
+            "6. ONLY return the description as a string, no extra text."
+        )
+
+        response_description = self.client.chat.completions.create(
+            model="gpt-4o-mini",
+            messages=[
+                {"role": "system", "content": "You are an AI assistant that generates infographic descriptions."},
+                {"role": "user", "content": description_prompt}
+            ]
+        )
+        generated_description = response_description.choices[0].message.content
+        generated_description = generated_description.strip() if generated_description else ""
+
+        return generated_title, generated_description
+
+def process(
+    input: str = None,
+    output: str = None,
+    input_data: Dict = None,
+    index_path: str = "faiss_infographics.index",
+    data_path: str = "infographics_data.npy",
+    topk: int = 7,
+    embed_model_path = "",
+    api_key: str="",
+    base_url: str=""
+) -> Union[bool, Dict]:
+    """
+    Process function for generating the title and subtitle for a single data object.
+
+    Args:
+        input (str, optional): Path to the input JSON file with a single data object.
+        output (str, optional): Path to the output JSON file.
+        input_data (Dict, optional): A single data dictionary (alternative to file input).
+        index_path (str, optional): Path to the FAISS index file.
+        data_path (str, optional): Path to the training data embeddings file.
+        topk (int, optional): Number of similar examples to retrieve.
+
+    Returns:
+        Union[bool, Dict]:
+          - If output is provided, returns True/False indicating success/failure.
+          - Otherwise, returns the updated data dictionary with generated titles.
+    """
+    try:
+        # Load the single data object
+        if input_data is None:
+            if input is None:
+                return False
+            with open(input, 'r', encoding='utf-8') as f:
+                data = json.load(f)
+        else:
+            data = input_data
+        if output is None:
+            output = input
+
+        if _has_required_titles(data):
+            if output:
+                with open(output, 'w', encoding='utf-8') as f:
+                    json.dump(data, f, ensure_ascii=False, indent=2)
+            return True
+
+        generator = _get_generator(
+            index_path=index_path,
+            data_path=data_path,
+            embed_model_path=embed_model_path,
+            api_key=api_key,
+            base_url=base_url
+        )
+
+        main_title, sub_title = generator.generate_title_description(data, topk=topk)
+
+        if "titles" not in data:
+            data["titles"] = {}
+        data["titles"]["main_title"] = main_title
+        data["titles"]["sub_title"] = sub_title
+
+        if output:
+            with open(output, 'w', encoding='utf-8') as f:
+                json.dump(data, f, ensure_ascii=False, indent=2)
+            return True
+
+        return True
+
+    except Exception as e:
+        print(f"Error in title generation: {str(e)}")
+        return False
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Title generator for chart data')
+    parser.add_argument('--input', type=str, required=True, help='Input JSON file path (single data object).')
+    parser.add_argument('--output', type=str, required=True, help='Output JSON file path.')
+    parser.add_argument('--index_path', type=str, default='faiss_infographics.index', help='FAISS index file path.')
+    parser.add_argument('--data_path', type=str, default='infographics_data.npy', help='Training data path.')
+    parser.add_argument('--topk', type=int, default=3, help='Number of similar examples to retrieve.')
+    parser.add_argument('--embed_model_path', type=str, default='', help='Sentence transformer path')
+    parser.add_argument('--api_key', type=str, default='', help='API key for LLM.')
+    parser.add_argument('--base_url', type=str, default='', help='Base URL for LLM.')
+
+    args = parser.parse_args()
+
+    success = process(
+        input=args.input,
+        output=args.output,
+        index_path=args.index_path,
+        data_path=args.data_path,
+        topk=args.topk,
+        embed_model_path=args.embed_model_path,
+        api_key=args.api_key,
+        base_url=args.base_url
+    )
+
+    if success:
+        print("Title generation completed successfully.")
+    else:
+        print("Title generation failed.")
+
+
+if __name__ == "__main__":
+    main()