!refactor: modularization

README.md

@@ -2,41 +2,21 @@
 
 ## Introduction
 
-Welcome to Auto Tabular Completion, a Python script designed to automatically fill in missing values in tabular data using advanced in-context learning techniques. This tool leverages the capabilities of LLMs to predict output values based on given input data.
+Welcome to Auto Tabular Completion. It automatically fill in missing values in tabular data using in-context learning techniques. This tool leverages the capabilities of LLMs to predict output values based on given input data.
 
-![Demo](./assets/demo.png)
+![Demo1](./assets/demo1.png)
 
 ## Quick Start
 
-1. Clone the repository:
-   ```bash
-   git clone https://github.com/Ki-Seki/autotab.git
-   cd autotab
-   ```
-2. Install the required dependencies:
-   ```bash
-   pip install -r requirements.txt
-   ```
-3. Running the Script
-    ```bash
-    python main.py \
-        --file <path_to_your_file>.xlsx \
-        --api_key <YOUR_OPENAI_API_KEY> \
-        [--base_url <BASE_URL>] \
-        [--model <MODEL_NAME>]
-    ```
+Check [demo1.ipynb](demo1.ipynb) and [demo2.ipynb](demo2.ipynb) for usage.
 
-## More
+> [!Note]
+> Give it a try with my personal key in the two demo notebooks (valid until 2024-11-01).
 
-### Optional Arguments
+## API
 
-- `--base_url`: Specify the base URL for the OpenAI API. This is useful if you're using a different API endpoint, such as one provided by [Silicon Flow](https://cloud.siliconflow.cn/). Default is the official OpenAI API endpoint.
-- `--model`: Choose the model to use for predictions. By default, the script uses `Qwen/Qwen2-7B-Instruct`.
+[Silicon Flow](https://cloud.siliconflow.cn/) provides free API.
 
-### Input Data Format
+## Input Format
 
-Your Excel file should contain columns prefixed with `[Input]` for input variables and `[Output]` for known or expected output values. Any columns not following this naming convention will be ignored during the prediction process.
-
-### Output
-
-After running the script, a new Excel file will be generated in the same directory as the input file, named `file_output.xlsx`, with all missing output values filled in based on the predictions made by the model.
+Your Excel file should contain columns prefixed with `[Input] ` for input variables and `[Output] ` for known or expected output values. Any columns not following this naming convention will be ignored during the prediction process.

autotab.py

@@ -1,126 +1,118 @@
-import argparse
-import concurrent.futures
-import os
 import re
 
 import openai
 import pandas as pd
-
-
-def load_excel(file_path):
-    """Load the Excel file and identify input and output columns."""
-    df = pd.read_excel(file_path)
-    input_columns = [col for col in df.columns if col.startswith("[Input]")]
-    output_columns = [col for col in df.columns if col.startswith("[Output]")]
-    return df, input_columns, output_columns
-
-
-def derive_in_context(data, input_columns, output_columns):
-    """Derive in-context learning data from the given DataFrame."""
-    n = len(data.dropna(subset=output_columns))
-    in_context = ""
-    for i in range(n):
-        in_context += "".join(
-            f"<{col.replace('[Input] ', '')}>{data[col].iloc[i]}</{col.replace('[Input] ', '')}>\n"
-            for col in input_columns
-        ) + "".join(
-            f"<{col.replace('[Output] ', '')}>{data[col].iloc[i]}</{col.replace('[Output] ', '')}>\n"
-            for col in output_columns
+from tqdm import tqdm
+
+
+class AutoTab:
+    def __init__(
+        self,
+        in_file_path: str,
+        out_file_path: str,
+        max_examples: int,
+        model_name: str,
+        api_key: str,
+        base_url: str,
+        generation_config: dict,
+        save_every: int,
+        instruction: str,
+    ):
+        self.in_file_path = in_file_path
+        self.out_file_path = out_file_path
+        self.max_examples = max_examples
+        self.model_name = model_name
+        self.api_key = api_key
+        self.base_url = base_url
+        self.generation_config = generation_config
+        self.save_every = save_every
+        self.instruction = instruction
+
+    # ─── IO ───────────────────────────────────────────────────────────────
+
+    def load_excel(self) -> tuple[pd.DataFrame, list, list]:
+        """Load the Excel file and identify input and output fields."""
+        df = pd.read_excel(self.in_file_path)
+        input_fields = [col for col in df.columns if col.startswith("[Input] ")]
+        output_fields = [col for col in df.columns if col.startswith("[Output] ")]
+        return df, input_fields, output_fields
+
+    # ─── LLM ──────────────────────────────────────────────────────────────
+
+    def openai_request(self, query: str) -> str:
+        """Make a request to an OpenAI-format API."""
+        client = openai.OpenAI(api_key=self.api_key, base_url=self.base_url)
+        response = client.chat.completions.create(
+            model=self.model_name,
+            messages=[{"role": "user", "content": query}],
+            **self.generation_config,
         )
-    return in_context, n
-
-
-def predict_output(in_context, input_data, input_columns, api_key, model, base_url):
-    """Predict the output values for the given input data using the OpenAI API."""
-    client = openai.OpenAI(api_key=api_key, base_url=base_url)
-    query = in_context + "".join(
-        f"<{col.replace('[Input] ', '')}>{input_data[col]}</{col.replace('[Input] ', '')}>\n"
-        for col in input_columns
-    )
-    response = client.chat.completions.create(
-        model=model, messages=[{"role": "user", "content": query}], max_tokens=64, n=1
-    )
-    predictions = response.choices[0].message.content.strip()
-    return predictions
-
-
-def extract_fields(prediction, output_columns):
-    """Extract fields from the prediction text based on output columns."""
-    extracted = {}
-    for col in output_columns:
-        tag = col.replace("[Output] ", "")
-        match = re.search(f"<{tag}>(.*?)</{tag}>", prediction)
-        extracted[tag] = match.group(1) if match else ""
-    return extracted
-
-
-def process_file(file_path, api_key, model, base_url):
-    """Process the entire Excel file and predict output values for missing entries."""
-    df, input_columns, output_columns = load_excel(file_path)
-    in_context, n = derive_in_context(df, input_columns, output_columns)
-
-    results = [None] * (len(df) - n)  # Pre-allocate a list for results
-    with concurrent.futures.ThreadPoolExecutor() as executor:
-        futures = {
-            executor.submit(
-                predict_output,
-                in_context,
-                df.iloc[i],
-                input_columns,
-                api_key,
-                model,
-                base_url,
-            ): i
-            - n
-            for i in range(n, len(df))
-        }
-        for future in concurrent.futures.as_completed(futures):
-            idx = futures[future]
-            try:
-                prediction = future.result()
-                extracted_fields = extract_fields(prediction, output_columns)
-                results[idx] = extracted_fields
-            except Exception as e:
-                results[idx] = {
-                    col.replace("[Output] ", ""): f"Error: {e}"
-                    for col in output_columns
-                }
-
-    # Assign results to the DataFrame
-    for i, extracted_fields in enumerate(results, start=n):
+        str_response = response.choices[0].message.content.strip()
+        return str_response
+
+    # ─── In-Context Learning ──────────────────────────────────────────────
+
+    def derive_incontext(
+        self, data: pd.DataFrame, input_columns: list[str], output_columns: list[str]
+    ) -> str:
+        """Derive the in-context prompt with angle brackets."""
+        n = min(self.max_examples, len(data.dropna(subset=output_columns)))
+        in_context = ""
+        for i in range(n):
+            in_context += "".join(
+                f"<{col.replace('[Input] ', '')}>{data[col].iloc[i]}</{col.replace('[Input] ', '')}>\n"
+                for col in input_columns
+            )
+            in_context += "".join(
+                f"<{col.replace('[Output] ', '')}>{data[col].iloc[i]}</{col.replace('[Output] ', '')}>\n"
+                for col in output_columns
+            )
+            in_context += "\n"
+        self.in_context = in_context
+        return in_context
+
+    def predict_output(
+        self, in_context: str, input_data: pd.DataFrame, input_fields: str
+    ):
+        """Predict the output values for the given input data using the API."""
+        query = (
+            self.instruction
+            + "\n\n"
+            + in_context
+            + "".join(
+                f"<{col.replace('[Input] ', '')}>{input_data[col]}</{col.replace('[Input] ', '')}>\n"
+                for col in input_fields
+            )
+        )
+        self.query_example = query
+        output = self.openai_request(query)
+        return output
+
+    def extract_fields(
+        self, response: str, output_columns: list[str]
+    ) -> dict[str, str]:
+        """Extract fields from the response text based on output columns."""
+        extracted = {}
         for col in output_columns:
-            tag = col.replace("[Output] ", "")
-            df.at[i, col] = extracted_fields.get(tag, "")
-
-    # Save the updated DataFrame to a new Excel file
-    output_file_path = os.path.splitext(file_path)[0] + "_output.xlsx"
-    df.to_excel(output_file_path, index=False)
-    print(f"Results saved to {output_file_path}")
-
-
-def main():
-    """Main function to parse arguments and execute the processing."""
-    parser = argparse.ArgumentParser(
-        description="Auto Tabular Completion with In-Context Learning"
-    )
-    parser.add_argument("--file", required=True, help="Path to the Excel file")
-    parser.add_argument("--api_key", required=True, help="OpenAI API key")
-    parser.add_argument(
-        "--base_url", required=False, help="OpenAI API base URL", default=None
-    )
-    parser.add_argument(
-        "--model",
-        required=False,
-        help="OpenAI model to use",
-        default="Qwen/Qwen2-7B-Instruct",
-    )
-
-    args = parser.parse_args(
-        "--file data/sample_nlp.xlsx --api_key sk-exhahhjfqyanmwewndukcqtrpegfdbwszkjucvcpajdufiah --base_url https://api.siliconflow.cn/v1".split()
-    )
-
-    process_file(args.file, args.api_key, args.model, args.base_url)
-
-
-if __name__ == "__main__":
-    main()
+            field = col.replace("[Output] ", "")
+            match = re.search(f"<{field}>(.*?)</{field}>", response)
+            extracted[col] = match.group(1) if match else ""
+        return extracted
+
+    # ─── Engine ───────────────────────────────────────────────────────────
+
+    def run(self):
+        data, input_fields, output_fields = self.load_excel()
+        in_context = self.derive_incontext(data, input_fields, output_fields)
+
+        num_existed_examples = len(data.dropna(subset=output_fields))
+
+        for i in tqdm(range(num_existed_examples, len(data))):
+            prediction = self.predict_output(in_context, data.iloc[i], input_fields)
+            extracted_fields = self.extract_fields(prediction, output_fields)
+            for field_name in output_fields:
+                data.at[i, field_name] = extracted_fields.get(field_name, "")
+            if i % self.save_every == 0:
+                data.to_excel(self.out_file_path, index=False)
+        data.to_excel(self.out_file_path, index=False)
+        print(f"Results saved to {self.out_file_path}")

demo1.ipynb

@@ -0,0 +1,110 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 3/3 [00:01<00:00,  2.60it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Results saved to data/en_qa_output.xlsx\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from autotab import AutoTab\n",
+    "\n",
+    "\n",
+    "autotab = AutoTab(\n",
+    "    in_file_path=\"data/en_qa_input.xlsx\",\n",
+    "    out_file_path=\"data/en_qa_output.xlsx\",\n",
+    "    max_examples=5,\n",
+    "    model_name=\"Qwen/Qwen2-7B-Instruct\",\n",
+    "    api_key=\"sk-exhahhjfqyanmwewndukcqtrpegfdbwszkjucvcpajdufiah\",\n",
+    "    base_url=\"https://public-beta-api.siliconflow.cn/v1\",\n",
+    "    generation_config={\"temperature\": 0, \"max_tokens\": 128},\n",
+    "    save_every=10,\n",
+    "    instruction='You should help me classify the questions and answer them.',\n",
+    ")\n",
+    "autotab.run()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "You should help me classify the questions and answer them.\n",
+      "\n",
+      "<Question>What is the capital of France?</Question>\n",
+      "<Category>Geography</Category>\n",
+      "<Answer>Paris</Answer>\n",
+      "\n",
+      "<Question>Who wrote '1984'?</Question>\n",
+      "<Category>Literature</Category>\n",
+      "<Answer>George Orwell</Answer>\n",
+      "\n",
+      "<Question>What is the largest planet in the solar system?</Question>\n",
+      "<Category>Astronomy</Category>\n",
+      "<Answer>Jupiter</Answer>\n",
+      "\n",
+      "<Question>Who painted the Mona Lisa?</Question>\n",
+      "<Category>Art</Category>\n",
+      "<Answer>Leonardo da Vinci</Answer>\n",
+      "\n",
+      "<Question>What is the currency of Japan?</Question>\n",
+      "<Category>Economics</Category>\n",
+      "<Answer>Yen</Answer>\n",
+      "\n",
+      "<Question>Who is the first president of the United States?</Question>\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(autotab.query_example)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "common",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

demo2.ipynb

@@ -0,0 +1,130 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 83/83 [01:19<00:00,  1.05it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Results saved to data/ch_patent_output.xlsx\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from autotab import AutoTab\n",
+    "\n",
+    "\n",
+    "instruction = \"\"\"你是一位专利分类专家，擅长判断专利是否属于智能化专利。请根据我提供的“智能化专利”定义、示例以及某专利的摘要，判断该专利是否属于智能化专利。\n",
+    "\n",
+    "## 智能化专利定义\n",
+    "\n",
+    "智能化专利是指通过利用先进智能技术（包括但不限于人工智能、高端芯片、量子信息、物联网、区块链、工业互联网和元宇宙）来提升系统、设备或流程的自主决策能力、自动化程度和优化效果的专利。这些专利旨在显著提高效率、减少错误、提升生产力，甚至实现新的功能和应用场景。例如，专利摘要中如果明确提到“使用人工智能技术”或“利用物联网技术”等词句的描述，则可以归为智能化专利。又例如，即使专利摘要中未出现这些明显词句，但如果描述了一种智能系统或设备，通过其自主运行或优化功能，能够在实际应用中减少人力操作，这类专利也可归为智能化专利。\n",
+    "\n",
+    "## 注意事项\n",
+    "\n",
+    "* 判断依据需要精简，一两句话即可。\n",
+    "* 请只回答判断依据和判断，不要包括任何多余内容。\n",
+    "\n",
+    "## 示例\"\"\"\n",
+    "\n",
+    "autotab = AutoTab(\n",
+    "    in_file_path=\"data/ch_patent_input.xlsx\",\n",
+    "    out_file_path=\"data/ch_patent_output.xlsx\",\n",
+    "    max_examples=4,\n",
+    "    model_name=\"Qwen/Qwen2-7B-Instruct\",\n",
+    "    api_key=\"sk-exhahhjfqyanmwewndukcqtrpegfdbwszkjucvcpajdufiah\",\n",
+    "    base_url=\"https://public-beta-api.siliconflow.cn/v1\",\n",
+    "    generation_config={\"temperature\": 0, \"max_tokens\": 512},\n",
+    "    save_every=10,\n",
+    "    instruction=instruction,\n",
+    ")\n",
+    "autotab.run()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "你是一位专利分类专家，擅长判断专利是否属于智能化专利。请根据我提供的“智能化专利”定义、示例以及某专利的摘要，判断该专利是否属于智能化专利。\n",
+      "\n",
+      "## 智能化专利定义\n",
+      "\n",
+      "智能化专利是指通过利用先进智能技术（包括但不限于人工智能、高端芯片、量子信息、物联网、区块链、工业互联网和元宇宙）来提升系统、设备或流程的自主决策能力、自动化程度和优化效果的专利。这些专利旨在显著提高效率、减少错误、提升生产力，甚至实现新的功能和应用场景。例如，专利摘要中如果明确提到“使用人工智能技术”或“利用物联网技术”等词句的描述，则可以归为智能化专利。又例如，即使专利摘要中未出现这些明显词句，但如果描述了一种智能系统或设备，通过其自主运行或优化功能，能够在实际应用中减少人力操作，这类专利也可归为智能化专利。\n",
+      "\n",
+      "## 注意事项\n",
+      "\n",
+      "* 判断依据需要精简，一两句话即可。\n",
+      "* 请只回答判断依据和判断，不要包括任何多余内容。\n",
+      "\n",
+      "## 示例\n",
+      "\n",
+      "<摘要>本发明公开了一种电动汽车智能节能充电桩系统及充电优化方法,属于汽车充电桩技术领域,包括车辆识别模块、充电预约模块、充电模块和服务器；所述充电预约模块用于电动汽车用户预约充电桩,建立预约通道,用户进行身份注册,将用户的注册信息进行储存,并生成识别标志；用户通过预约通道预约充电桩,设置充电预约时间YT、等待计费价格JG和等待时长DT,并将预约信息发送到车辆识别模块；所述车辆识别模块用于识别电动车辆,当识别成功时,通过充电模块对电动车辆进行充电；通过设置充电预约时间YT、等待计费价格JG和等待时长DT,督促用户在预约后能够按时到达停车位,避免预约后,也没有按时充电,造成资源的浪费。</摘要>\n",
+      "<判断依据>此发明利用了电动汽车充电预约系统，有智能优化方法，但未明确提及人工智能、物联网、量子信息等其他智能技术，主要依赖于电动汽车充电设备的智能化管理，因此可以归类为智能化专利。</判断依据>\n",
+      "<判断>是</判断>\n",
+      "\n",
+      "<摘要>本发明公开了一种超声成像方法及系统、存储介质、处理器和计算机设备。其中,该方法包括：获取胎儿头颅的三维体数据,其中,三维体数据是经超声对头颅进行扫查后得到的数据；分解三维体数据,生成预定数量的二维切面图像；分别对预定数量的二维切面图像进行分割,得到二维切面图像中颅内区域的轮廓；根据预定数量的二维切面图像中颅内区域的轮廓,拟合出头颅的三维颅脑轮廓；根据三维颅脑轮廓,确定头颅中颅脑的体积。本发明解决了相关技术中的测量方式,准确率较低,测量速度较慢的技术问题。</摘要>\n",
+      "<判断依据>此发明利用了超声成像技术，通过三维体数据处理和图像分割等方法，实现了对胎儿头颅的精确测量，体现了智能化处理和优化效果，属于智能化专利范畴。</判断依据>\n",
+      "<判断>是</判断>\n",
+      "\n",
+      "<摘要>一种锂离子电池正极活性物质磷酸亚铁锂的制备方法,该方法包括将高\\r\\n\\r\\n\\r\\n\\r\\n分子聚合物、磷源化合物、铁源化合物和锂源化合物混合后进行烧结,其中,\\r\\n\\r\\n\\r\\n\\r\\n所述铁源化合物为铁的非二价化合物；所述高分子聚合物、磷源化合物和锂\\r\\n\\r\\n\\r\\n\\r\\n源化合物与铁的非二价化合物的混合在溶液中进行,得到的混合溶液除去溶\\r\\n\\r\\n\\r\\n\\r\\n剂之后得到凝胶体,然后将所得凝胶体进行烧结。本发明选用铁的非二价化\\r\\n\\r\\n\\r\\n\\r\\n合物为原料,大大降低成本,而且制备方法简单易行。利用本发明方法制得\\r\\n\\r\\n\\r\\n\\r\\n的磷酸亚铁锂作为正极活性物质,使得电池具有较高的首次放电比容量和良\\r\\n\\r\\n\\r\\n\\r\\n好的循环性能。并且用本发明方法制得的磷酸亚铁锂成分固定、颗粒均匀、\\r\\n\\r\\n\\r\\n\\r\\n导电性好。</摘要>\n",
+      "<判断依据>此发明主要涉及电池制备方法，未明确提及使用了先进智能技术，如人工智能、物联网等。</判断依据>\n",
+      "<判断>否</判断>\n",
+      "\n",
+      "<摘要>本实用新型公开了一种水泵井壁支撑装置,用于实现水泵与井壁之间的分隔支撑,解决了现有类似装置的结构复杂,使用时不能保证顺畅的问题,所采取的技术方案：一种水泵井壁支撑装置,用于实现水泵与井壁之间的支撑,包括支撑件,支撑件包括用于联接在水泵壳体上的联接部,其特征在于,联接部上设有若干支撑部,这些支撑部在联接部的外周面上突出设置,这些支撑部呈放射状布置在联接部上。本水泵井壁支撑装置结构简单,造价低,使用稳定性好,能够为水泵与井壁之间提供良好的分隔支撑作用。  来自马克数据网</摘要>\n",
+      "<判断依据>此发明主要关注于水泵井壁支撑装置的结构设计，未提及使用任何智能技术，如人工智能、物联网等，因此主要依赖于物理结构的优化，而非智能技术提升效率或自动化。</判断依据>\n",
+      "<判断>否</判断>\n",
+      "\n",
+      "<摘要>本实用新型提供了一种翻转梁,包括底座、前盖,底座上设有导向结构,导向结构包括导向件、弹簧,导向件的一端从底座的上端对应开孔穿出,弹簧的一端连接于导向件,弹簧的另一端连接于底座,底座设有一对用于限制导向件直线移动的第一限位件。本实用新型提供的翻转梁及冰箱,其连接于底座与导向件之间的弹簧使得导向件可以伸缩,配合设于底座、前盖和导向件上的限位件,从而让导向凸起进入导向槽中并向下运动,待导向凸起脱离导向槽并进入导向凸起安装座中时利用弹簧的回复作用使导向件回到初始位置,同时在导向结构沿导向槽滑动的过程中起到缓冲的作用,并方便了弹簧的安装,同时导向件大部分结构位于翻转梁内部,开门时不影响冰箱的外观。</摘要>\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(autotab.query_example)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "common",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}