| import torch
|
| import torch.nn.functional as F
|
| import os
|
| from model.config import ModelConfig
|
| from model.transformer import GPT
|
| from model.tokenizer import AdvancedTokenizer
|
|
|
| class InferenceEngine:
|
| def __init__(self, model_path='sail.pt'):
|
| self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
| self.draft_model = None
|
|
|
| if not os.path.exists(model_path):
|
| print(f"Warning: {model_path} not found. Agent will be uninitialized.")
|
| self.model = None
|
| return
|
|
|
| print(f"Loading model from {model_path}...")
|
| try:
|
| checkpoint = torch.load(model_path, map_location=self.device, weights_only=False)
|
|
|
| if isinstance(checkpoint, dict) and 'config' in checkpoint:
|
| self.config = checkpoint['config']
|
| self.config.device = self.device
|
| self.tokenizer = AdvancedTokenizer(vocab_size=self.config.vocab_size)
|
| if 'vocab' in checkpoint:
|
| self.tokenizer.word_to_id = checkpoint['vocab']
|
| self.tokenizer.id_to_word = {v: k for k, v in self.tokenizer.word_to_id.items()}
|
| state_dict = checkpoint['model_state_dict']
|
|
|
| state_dict = {k.replace('_orig_mod.', ''): v for k, v in state_dict.items()}
|
| self.model = GPT(self.config).to(self.device)
|
| self.model.load_state_dict(state_dict)
|
| else:
|
| import pickle
|
| with open('tokenizer.pkl', 'rb') as f:
|
| self.tokenizer = pickle.load(f)
|
| self.config = ModelConfig()
|
| self.model = GPT(self.config).to(self.device)
|
| self.model.load_state_dict(checkpoint)
|
|
|
| self.model.eval()
|
| print("Model loaded successfully.")
|
| except Exception as e:
|
| print(f"Error loading model: {e}")
|
| self.model = None
|
|
|
| def speculative_generate(self, idx, max_new_tokens, K=4):
|
| """
|
| Speculative Decoding: Generate multiple tokens using a draft model,
|
| then verify them in one pass with the big model.
|
| """
|
| if self.draft_model is None:
|
| return self.model.generate(idx, max_new_tokens)
|
|
|
| generated = 0
|
| while generated < max_new_tokens:
|
| T = idx.shape[1]
|
| draft_idx = idx.clone()
|
| for _ in range(min(K, max_new_tokens - generated)):
|
| logits, _ = self.draft_model(draft_idx)
|
| next_token = torch.multinomial(F.softmax(logits[:, -1, :], dim=-1), 1)
|
| draft_idx = torch.cat([draft_idx, next_token], dim=1)
|
|
|
|
|
| full_logits, _ = self.model(draft_idx)
|
|
|
|
|
| idx = draft_idx
|
| generated += K
|
| if generated >= max_new_tokens: break
|
|
|
| return idx
|
|
|
| def chat(self, prompt, system_prompt="You are a smart, agentic AI assistant.", max_steps=3):
|
| if not self.model: return "Error: Model not loaded."
|
|
|
| from agent.tool_executor import parse_and_execute_tools
|
| current_text = f"[SYSTEM] {system_prompt} [USER] {prompt} [THOUGHT]"
|
|
|
| for step in range(max_steps):
|
| token_ids = self.tokenizer.encode(current_text)
|
| idx = torch.tensor(token_ids, dtype=torch.long).unsqueeze(0).to(self.device)
|
|
|
| with torch.no_grad():
|
| out_idx = self.speculative_generate(idx, max_new_tokens=256)
|
|
|
| generated_text = self.tokenizer.decode(out_idx[0].tolist())
|
|
|
|
|
| tool_result = parse_and_execute_tools(generated_text)
|
| if tool_result:
|
| current_text = generated_text + " " + tool_result
|
| continue
|
| else:
|
|
|
| if "[THOUGHT]" in generated_text and "error" in generated_text.lower():
|
| current_text = generated_text + " [THOUGHT] I detected a potential error. I must correct it."
|
| continue
|
| return generated_text
|
| return generated_text
|
|
|
| def generate(self, prompt, max_new_tokens=200, **kwargs):
|
| if not self.model: return "Model not loaded."
|
| token_ids = self.tokenizer.encode(prompt)
|
| idx = torch.tensor(token_ids, dtype=torch.long).unsqueeze(0).to(self.device)
|
| with torch.no_grad():
|
| out_idx = self.model.generate(idx, max_new_tokens, **kwargs)
|
| return self.tokenizer.decode(out_idx[0].tolist())
|
|
|
