Upload chat.py with huggingface_hub

chat.py

@@ -0,0 +1,204 @@
+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)