kofdai commited on
Commit
e117789
·
verified ·
1 Parent(s): 30db877

Upload src/iath_encoder.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. src/iath_encoder.py +168 -0
src/iath_encoder.py ADDED
@@ -0,0 +1,168 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import struct
2
+ import zstandard as zstd
3
+ from datetime import datetime
4
+ import json # これを追加
5
+
6
+ class IathEncoder:
7
+ """
8
+ Knowledge Tileオブジェクトを.iath互換の圧縮バイナリにエンコードします。
9
+ """
10
+
11
+ def _encode_reviewer_reference(self, reviewer: dict) -> bytes:
12
+ """
13
+ レビュアー情報をエンコードします。
14
+ 当面はダミー実装とし、レビュアーIDを固定長で返します。
15
+ 将来的にはVerifier Dictionaryを参照するインデックスを返す必要があります。
16
+ """
17
+ reviewer_id = reviewer.get("reviewer_id", "unknown").encode('utf-8')
18
+ return struct.pack("<36s", reviewer_id[:36]) # UUID string length
19
+
20
+ def _encode_string(self, s: str) -> bytes:
21
+ """NULL終端のUTF-8文字列をエンコードします。"""
22
+ return s.encode('utf-8') + b'\0'
23
+
24
+ def _encode_metadata(self, metadata: dict) -> bytes:
25
+ """メタデータをバイナリ化します。"""
26
+ kid = self._encode_string(metadata["knowledge_id"])
27
+ topic = self._encode_string(metadata["topic"])
28
+ created_at_iso = metadata.get("created_at", datetime.now().isoformat())
29
+ created_at = created_at_iso.encode('ascii')[:27] # ISO format with Z
30
+ return kid + topic + created_at
31
+
32
+ def _encode_coordinates(self, coordinates: dict) -> bytes:
33
+ """座標をバイナリ化(6つの浮動小数点数)。"""
34
+ medical_space = coordinates["medical_space"]
35
+ meta_space = coordinates["meta_space"]
36
+
37
+ return struct.pack(
38
+ "<ffffff",
39
+ float(medical_space[0]), float(medical_space[1]), float(medical_space[2]),
40
+ float(meta_space[0]), float(meta_space[1]), float(meta_space[2])
41
+ )
42
+
43
+ def _encode_content(self, content: dict) -> bytes:
44
+ """コンテンツ(テキスト)をバイナリ化します。"""
45
+ thinking = content["thinking_process"].encode('utf-8')
46
+ response = content["final_response"].encode('utf-8')
47
+
48
+ # 各パートの長さを前に付けて連結
49
+ result = struct.pack("<I", len(thinking)) + thinking
50
+ result += struct.pack("<I", len(response)) + response
51
+ return result
52
+
53
+ def _encode_verification(self, verification: dict) -> bytes:
54
+ """検証履歴をバイナリ化します。"""
55
+ status_map = {
56
+ "pending_review": 0, "partial_verified": 1,
57
+ "verified": 2, "expert_confirmed": 3
58
+ }
59
+ status_code = status_map.get(verification.get("status", "pending_review"), 0)
60
+
61
+ initial_certainty = int(verification.get("initial_certainty", 0))
62
+ reviewer_count = len(verification.get("reviewers", []))
63
+
64
+ result = struct.pack("<BBI", status_code, initial_certainty, reviewer_count)
65
+
66
+ for reviewer in verification.get("reviewers", []):
67
+ result += self._encode_reviewer_reference(reviewer)
68
+
69
+ return result
70
+
71
+ def encode_tile(self, tile: dict) -> bytes:
72
+ """
73
+ 単一のKnowledge Tileをエンコードし、zstdで圧縮します。
74
+
75
+ Args:
76
+ tile (dict): Knowledge Tileオブジェクト。
77
+
78
+ Returns:
79
+ bytes: 圧縮されたバイナリデータ。
80
+ """
81
+ # 各セクションをエンコード
82
+ metadata_bin = self._encode_metadata(tile["metadata"])
83
+ coord_bin = self._encode_coordinates(tile["coordinates"])
84
+ content_bin = self._encode_content(tile["content"])
85
+ verification_bin = self._encode_verification(tile["verification"])
86
+
87
+ # NOTE: reasoning_path, source, historyなどは今回省略し、主要な部分のみ実装
88
+
89
+ # 長さプレフィックスを付けて連結
90
+ uncompressed = b"".join([
91
+ struct.pack("<I", len(metadata_bin)), metadata_bin,
92
+ struct.pack("<I", len(coord_bin)), coord_bin,
93
+ struct.pack("<I", len(content_bin)), content_bin,
94
+ struct.pack("<I", len(verification_bin)), verification_bin,
95
+ ])
96
+
97
+ # zstdで圧縮
98
+ cctx = zstd.ZstdCompressor(level=19)
99
+ compressed = cctx.compress(uncompressed)
100
+
101
+ return compressed
102
+
103
+ def encode_batch(self, tiles: List[Dict], domain_code: int = 1) -> bytes:
104
+ """
105
+ 複数の知識タイルを受け取り、完全な.iathデータベースファイルのバイナリを生成します。
106
+
107
+ Args:
108
+ tiles (List[Dict]): エンコードする知識タイルの辞書のリスト。
109
+ domain_code (int): ヘッダーに書き込むドメインコード (1: medical, 2: legal, etc.)。
110
+
111
+ Returns:
112
+ bytes: 完全な.iathファイルのバイナリコンテンツ。
113
+ """
114
+ print(f"--- {len(tiles)}件のタイルのバッチエンコード開始 (ドメインコード: {domain_code}) ---")
115
+
116
+ index = []
117
+ data_chunks = []
118
+ current_offset = 0
119
+
120
+ # 1. 各タイルを個別にエンコードし、データチャンクとインデックスを作成
121
+ for tile in tiles:
122
+ tile_id = tile.get("metadata", {}).get("knowledge_id")
123
+ if not tile_id:
124
+ print("警告: knowledge_idのないタイルをスキップします。")
125
+ continue
126
+
127
+ compressed_data = self.encode_tile(tile)
128
+ data_length = len(compressed_data)
129
+
130
+ index.append({"id": tile_id, "offset": current_offset, "length": data_length})
131
+ data_chunks.append(compressed_data)
132
+
133
+ current_offset += data_length
134
+
135
+ print(" - 全タイルの個別エンコード完了。")
136
+
137
+ # 2. インデックスセクションをシリアライズ
138
+ index_binary = json.dumps(index, ensure_ascii=False).encode('utf-8')
139
+ print(f" - インデックス作成完了 (サイズ: {len(index_binary)} bytes)")
140
+
141
+ # 3. データセクションを結合
142
+ data_section = b"".join(data_chunks)
143
+
144
+ # 4. ヘッダーを作成
145
+ header_size = 64
146
+ index_offset = header_size
147
+ data_offset = index_offset + len(index_binary)
148
+
149
+ checksum = b'\0' * 32
150
+
151
+ header = struct.pack(
152
+ "<4sIBB32sQQ6x",
153
+ b'ILMA', # Magic number
154
+ 1, # Version
155
+ domain_code, # ドメインコードを引数から設定
156
+ 1, # Compression Type (0x01=zstd)
157
+ checksum,
158
+ index_offset,
159
+ data_offset
160
+ )
161
+ print(" - ヘッダー作成完了。")
162
+
163
+ # 5. すべてのセクションを結合
164
+ full_db_content = header + index_binary + data_section
165
+ print("--- バッチエンコード完了 ---")
166
+
167
+ return full_db_content
168
+