Create AppleAI-converter-colab.py

AppleAI-converter-colab.py

@@ -0,0 +1,74 @@
+import os
+os.system('pip install mlx')
+import mlx as mx
+import mlx.nn as mx_nn
+import torch
+import torch.nn as nn
+import numpy as np
+
+device = torch.device('cuda')
+
+CONFIG = {
+    "model_path": "NeoChess/chessy_model.pth",
+    "backup_model_path": "NeoChess/chessy_modelt-1.pth",
+}
+
+class NN1(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.embedding = nn.Embedding(13, 64)
+        self.attention = nn.MultiheadAttention(embed_dim=64, num_heads=16)
+        self.neu = 512
+        self.neurons = nn.Sequential(
+            nn.Linear(4096, self.neu),
+            nn.ReLU(),
+            nn.Linear(self.neu, self.neu), 
+            nn.ReLU(),
+            nn.Linear(self.neu, self.neu), 
+            nn.ReLU(),
+            nn.Linear(self.neu, self.neu), 
+            nn.ReLU(),
+            nn.Linear(self.neu, self.neu), 
+            nn.ReLU(),
+            nn.Linear(self.neu, self.neu), 
+            nn.ReLU(),
+            nn.Linear(self.neu, self.neu), 
+            nn.ReLU(),
+            nn.Linear(self.neu, self.neu), 
+            nn.ReLU(),
+            nn.Linear(self.neu, self.neu), 
+            nn.ReLU(),
+            nn.Linear(self.neu, self.neu), 
+            nn.ReLU(),
+            nn.Linear(self.neu, self.neu), 
+            nn.ReLU(),
+            nn.Linear(self.neu, self.neu), 
+            nn.ReLU(),
+            nn.Linear(self.neu, self.neu), 
+            nn.ReLU(),
+            nn.Linear(self.neu, 64), 
+            nn.ReLU(),
+            nn.Linear(64, 4)
+        )
+
+    def forward(self, x):
+        x = self.embedding(x)
+        x = x.permute(1, 0, 2)
+        attn_output, _ = self.attention(x, x, x)
+        x = attn_output.permute(1, 0, 2).contiguous()
+        x = x.view(x.size(0), -1)
+        x = self.neurons(x)
+        return x
+
+model = NN1().to(device)
+try:
+    model.load_state_dict(torch.load(CONFIG['model_path'], map_location=device))
+    print(f"Loaded model from {CONFIG['model_path']}")
+except FileNotFoundError:
+    try:
+        model.load_state_dict(torch.load(CONFIG["backup_model_path"], map_location=device))
+        print(f"Loaded backup model from {CONFIG['backup_model_path']}")
+    except FileNotFoundError:
+        print("No model file found, starting from scratch.")
+weights = {k: v.detach().cpu().numpy() for k, v in model.state_dict().items()}
+np.savez("chessy_model_mlx.npz", **weights)