File size: 4,139 Bytes
e5b79b7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | import json
import os
import torch
from torch.utils.data import Dataset
class InstructionDataset(Dataset):
"""
Dataset for Instruction Tuning (Q&A).
Format: [INSTRUCTION] Question [RESPONSE] Answer <EOS>
"""
def __init__(self, data_path, tokenizer, block_size):
self.tokenizer = tokenizer
self.block_size = block_size
self.data = []
try:
import sys
import os
# Ensure the Rust .dll/.so is on path
core_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'agent_core', 'target', 'release')
if core_path not in sys.path:
sys.path.append(core_path)
import agent_core
# Check if the attribute actually exists before calling to prevent unhandled AttributeError
if hasattr(agent_core, 'parse_instruction_dataset_safe'):
self.data = agent_core.parse_instruction_dataset_safe(data_path)
print(f"Loaded {len(self.data)} instruction pairs via Rust Memory-Safe Core.")
else:
raise AttributeError(f"module 'agent_core' has no attribute 'parse_instruction_dataset_safe'")
except (ImportError, AttributeError, Exception) as e:
# The Rust extension handles all JSON loading, error checking,
# and parallel string formatting in memory-safe C-bindings natively
print(f"Warning: Rust `agent_core` fallback triggered. Reason: {e}")
print("Falling back to native Python parsing.")
self._fallback_load(data_path)
def _fallback_load(self, data_path):
try:
import json
with open(data_path, 'r', encoding='utf-8') as f:
raw_data = json.load(f)
for item in raw_data:
if 'instruction' in item and 'response' in item:
system_prompt = item.get('system', 'You are a helpful, smart AI assistant.')
thought = item.get('thought', '')
tools = item.get('tools', [])
thought_str = f"[THOUGHT] {thought} " if thought else ""
tools_str = ""
for tool in tools:
tools_str += f"[TOOL_CALL] {tool.get('name')} [TOOL_ARG] {tool.get('args')} [TOOL_RESULT] {tool.get('result')} "
text = f"[SYSTEM] {system_prompt} [USER] {item['instruction']} {thought_str}{tools_str}[RESPONSE] {item['response']} <EOS>"
self.data.append(text)
except Exception as e:
print(f"Error loading instruction data: {e}")
self.data = []
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
text = self.data[idx]
# Tokenize
# Note: We need a custom encode method that handles the special tags if they aren't in vocab
# For now, we assume the tokenizer splits them or we add them as special tokens
# Simple encoding for now, assuming tokenizer handles it or splits by space
# Ideally, we should add [INSTRUCTION] and [RESPONSE] to tokenizer specials
# For simplicity in this iteration, we'll just encode the text string
# The tokenizer.encode needs to be robust
# We will use the tokenizer's encode method
# If tokenizer doesn't have the special tokens, it might split them.
# We should add them to the tokenizer in train.py
ids = self.tokenizer.encode(text)
# Pad or Truncate
if len(ids) > self.block_size:
ids = ids[:self.block_size]
else:
ids = ids + [self.tokenizer.word_to_id.get("<PAD>", 0)] * (self.block_size - len(ids))
x = torch.tensor(ids[:-1], dtype=torch.long)
y = torch.tensor(ids[1:], dtype=torch.long)
return x, y
|