| import torch |
| import sys |
| import os |
|
|
| def run_nexus(weights_path): |
| sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src')) |
|
|
| from src.trainer import load_nexus |
| from tokenizers import Tokenizer |
| import torch.nn.functional as F |
|
|
| model, config = load_nexus(weights_path) |
|
|
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
| model = model.to(device) |
|
|
| tokenizer_path = os.path.join(os.path.dirname(weights_path), '..', 'data', 'tokenizer.json') |
| tokenizer_path = os.path.normpath(tokenizer_path) |
|
|
| if not os.path.exists(tokenizer_path): |
| tokenizer_path = os.path.join(os.path.dirname(__file__), 'data', 'tokenizer.json') |
|
|
| tokenizer = Tokenizer.from_file(tokenizer_path) |
|
|
| print("\n{Nexus SmAll v1} Chat Interface") |
| print("{Nexus SmAll v1} Type 'exit' to quit, 'clear' to reset conversation") |
| print("{Nexus SmAll v1} Type '--temp 0.5' to change temperature") |
| print("{Nexus SmAll v1} Type '--help' for all commands\n") |
|
|
| bos_id = tokenizer.token_to_id("<bos>") if tokenizer.token_to_id("<bos>") is not None else 1 |
| eos_id = tokenizer.token_to_id("<eos>") if tokenizer.token_to_id("<eos>") is not None else 2 |
|
|
| conversation = [bos_id] |
|
|
| temperature = 0.2 |
| top_k = 40 |
| top_p = 0.9 |
| max_tokens = 128 |
| repetition_penalty = 1.2 |
|
|
| while True: |
| try: |
| user_input = input("You: ").strip() |
|
|
| if not user_input: |
| continue |
| if user_input.lower() == 'exit': |
| print("Goodbye!") |
| break |
| elif user_input.lower() == 'clear': |
| conversation = [bos_id] |
| print("[Conversation reset]") |
| continue |
| elif user_input.startswith('--'): |
| parts = user_input.split() |
| if parts[0] == '--temp' and len(parts) >= 2: |
| temperature = float(parts[1]) |
| print(f"[temperature={temperature}]") |
| continue |
| elif parts[0] == '--help': |
| print("Commands:") |
| print(" --temp <value> Set temperature (default 0.2)") |
| print(" --topk <value> Set top_k (default 40)") |
| print(" --topp <value> Set top_p (default 0.9)") |
| print(" --tokens <value> Set max new tokens (default 128)") |
| print(" --rep <value> Set repetition penalty (default 1.2)") |
| print(" clear Reset conversation") |
| print(" exit Exit") |
| continue |
| elif parts[0] == '--topk' and len(parts) >= 2: |
| top_k = int(parts[1]) |
| print(f"[top_k={top_k}]") |
| continue |
| elif parts[0] == '--topp' and len(parts) >= 2: |
| top_p = float(parts[1]) |
| print(f"[top_p={top_p}]") |
| continue |
| elif parts[0] == '--tokens' and len(parts) >= 2: |
| max_tokens = int(parts[1]) |
| print(f"[max_tokens={max_tokens}]") |
| continue |
| elif parts[0] == '--rep' and len(parts) >= 2: |
| repetition_penalty = float(parts[1]) |
| print(f"[repetition_penalty={repetition_penalty}]") |
| continue |
|
|
| prompt = f"\nUser: {user_input}\nAssistant:" |
| prompt_ids = tokenizer.encode(prompt).ids |
| input_ids = conversation + prompt_ids |
|
|
| if len(input_ids) > config.max_seq_len: |
| input_ids = input_ids[-config.max_seq_len + 64:] |
|
|
| input_tensor = torch.tensor([input_ids], dtype=torch.long, device=device) |
|
|
| generated_ids, full_ids = _generate_with_rep_penalty( |
| model, input_tensor, max_new_tokens=max_tokens, |
| temperature=temperature, top_k=top_k, top_p=top_p, |
| repetition_penalty=repetition_penalty, |
| eos_id=eos_id, |
| ) |
|
|
| response_ids = full_ids[0, input_tensor.shape[1]:].tolist() |
| response_text = tokenizer.decode(response_ids) |
|
|
| if "<eos>" in response_text: |
| response_text = response_text[:response_text.index("<eos>")] |
| if "<bos>" in response_text: |
| response_text = response_text.replace("<bos>", "") |
| if "User:" in response_text: |
| response_text = response_text[:response_text.index("User:")] |
| if "Assistant:" in response_text: |
| response_text = response_text.replace("Assistant:", "") |
|
|
| response_text = response_text.strip() |
|
|
| if len(response_text) < 2: |
| response_text = "[no response]" |
|
|
| print(f"Nexus SmAll v1: {response_text}") |
|
|
| conversation = full_ids[0].tolist() |
| if eos_id is not None: |
| conversation.append(eos_id) |
|
|
| except KeyboardInterrupt: |
| print("\nGoodbye!") |
| break |
| except Exception as e: |
| print(f"[Error] {e}") |
| continue |
|
|
| def _generate_with_rep_penalty(model, input_ids, max_new_tokens, temperature, top_k, top_p, repetition_penalty, eos_id): |
| model.eval() |
|
|
| for _ in range(max_new_tokens): |
| seq_len = input_ids.shape[1] |
| if seq_len > model.config.max_seq_len: |
| input_ids = input_ids[:, -model.config.max_seq_len:] |
|
|
| with torch.no_grad(): |
| logits = model(input_ids, 0) |
| logits = logits[:, -1, :] |
|
|
| if repetition_penalty != 1.0: |
| for batch_idx in range(logits.shape[0]): |
| for token_idx in range(input_ids.shape[1]): |
| token = input_ids[batch_idx, token_idx].item() |
| if logits[batch_idx, token] < 0: |
| logits[batch_idx, token] *= repetition_penalty |
| else: |
| logits[batch_idx, token] /= repetition_penalty |
|
|
| logits = logits / temperature |
|
|
| if top_k > 0: |
| top_k_values, _ = torch.topk(logits, min(top_k, logits.size(-1))) |
| min_top_k = top_k_values[:, -1].unsqueeze(-1) |
| logits = torch.where(logits < min_top_k, |
| torch.full_like(logits, float('-inf')), logits) |
|
|
| if top_p > 0 and top_p < 1.0: |
| sorted_logits, sorted_indices = torch.sort(logits, descending=True) |
| cumulative_probs = torch.cumsum(torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1) |
| sorted_indices_to_remove = cumulative_probs > top_p |
| sorted_indices_to_remove[:, 0] = False |
| indices_to_remove = torch.zeros_like(logits, dtype=torch.bool) |
| indices_to_remove = indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove) |
| logits = torch.where(indices_to_remove, |
| torch.full_like(logits, float('-inf')), logits) |
|
|
| probs = torch.nn.functional.softmax(logits, dim=-1) |
| next_token = torch.multinomial(probs, num_samples=1) |
|
|
| input_ids = torch.cat([input_ids, next_token], dim=-1) |
|
|
| if eos_id is not None and next_token.item() == eos_id: |
| break |
|
|
| return None, input_ids |
|
|
| if __name__ == "__main__": |
| import argparse |
|
|
| parser = argparse.ArgumentParser(description="Nexus SmAll v1 Chat") |
| parser.add_argument("--weights", type=str, default="weights/nexus_final.pt", |
| help="Path to model weights (.pt file)") |
| parser.add_argument("--temp", type=float, default=0.2, |
| help="Temperature (default: 0.2)") |
| parser.add_argument("--top_k", type=int, default=40, |
| help="Top-k sampling (default: 40)") |
| parser.add_argument("--top_p", type=float, default=0.9, |
| help="Top-p sampling (default: 0.9)") |
| parser.add_argument("--max_tokens", type=int, default=128, |
| help="Max new tokens (default: 128)") |
| args = parser.parse_args() |
|
|
| if not os.path.exists(args.weights): |
| print(f"[Error] Weights not found: {args.weights}") |
| print("Make sure training completed successfully.") |
| input("Press Enter to exit...") |
| sys.exit(1) |
|
|
| run_nexus(args.weights) |
