Add full SocrateX API directly on SOCRATE class (create_config, new, make_trainer, generate_data, load_data)

model.py

@@ -296,17 +296,151 @@ class SOCRATE(PreTrainedModel):
             param.requires_grad = True
         print("Encoder has been unfrozen.")
         
-    def summary(self):
+    # ──────────────────────────────────────────────────────────────────────────
+    # Class-level API — everything SocrateX can do, directly on the model
+    # ──────────────────────────────────────────────────────────────────────────
+
+    @staticmethod
+    def create_config(
+        d_model=256,
+        nhead=4,
+        num_layers=4,
+        dim_feedforward=1024,
+        activation="gelu",
+        norm_first=True,
+        max_len=512,
+        pool_height=4,
+    ):
         """
-        Prints a summary of the model's parameters.
+        Create a custom architecture config without needing to import SocrateX separately.
+
+        Example::
+
+            model = AutoModel.from_pretrained("ihatebaselines/Socrate", trust_remote_code=True)
+            cfg = model.create_config(d_model=512, nhead=8, num_layers=6, dim_feedforward=2048)
+            tok = model.make_tokenizer()
+            new_model = model.new(config=cfg, tokenizer=tok)
+        """
+        from configuration_socrate import SocrateConfig
+        return SocrateConfig(
+            d_model=d_model,
+            nhead=nhead,
+            num_layers=num_layers,
+            dim_feedforward=dim_feedforward,
+            activation=activation,
+            norm_first=norm_first,
+            max_len=max_len,
+            pool_height=pool_height,
+        )
+
+    @staticmethod
+    def make_tokenizer(path=None):
+        """
+        Initialize a fresh BPE tokenizer from scratch, or load one from a file.
+
+        Example::
+
+            tok = model.make_tokenizer()                         # fresh tokenizer
+            tok = model.make_tokenizer("ocr_bpe_tokenizer.json") # load from file
+        """
+        if path is not None:
+            from tokenizers import Tokenizer
+            return Tokenizer.from_file(path)
+        try:
+            from tokenizer import init_tokenizer
+        except ImportError:
+            from SocrateX.tokenizer import init_tokenizer
+        return init_tokenizer()
+
+    @classmethod
+    def new(cls, config, tokenizer, device="cuda"):
+        """
+        Build a brand-new SOCRATE model from a config + tokenizer.
+        No pretrained weights — starts from scratch.
+
+        Example::
+
+            cfg = model.create_config(d_model=256, nhead=4, num_layers=4, dim_feedforward=1024)
+            tok = model.make_tokenizer()
+            my_model = model.new(config=cfg, tokenizer=tok, device="cpu")
+            print(my_model.summary())
+        """
+        import torch
+        hf_config = cls.create_config(
+            d_model=config.d_model if hasattr(config, 'd_model') else 256,
+            nhead=config.nhead if hasattr(config, 'nhead') else 4,
+            num_layers=config.num_layers if hasattr(config, 'num_layers') else 4,
+            dim_feedforward=config.dim_feedforward if hasattr(config, 'dim_feedforward') else 1024,
+        )
+        hf_config.vocab_size = tokenizer.get_vocab_size()
+        hf_config.pad_id = tokenizer.token_to_id("<pad>")
+        hf_config.bos_id = tokenizer.token_to_id("<bos>")
+        hf_config.eos_id = tokenizer.token_to_id("<eos>")
+        return cls(hf_config, tokenizer=tokenizer, sx_config=config).to(device)
+
+    def make_trainer(self, dataloader, optimizer, criterion, device=None):
         """
-        total_params = sum(p.numel() for p in self.parameters())
-        trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
-        print(f"=== SOCRATE Model Summary ===")
-        print(f"Total parameters: {total_params / 1e6:.2f} M")
-        print(f"Trainable parameters: {trainable_params / 1e6:.2f} M")
-        print(f"d_model: {self.d_model}")
-        print("=============================")
+        Returns a Trainer object wired to this model.
+
+        Example::
+
+            loader = model.make_dataset(images, labels).to_loader(batch_size=16)
+            opt    = torch.optim.AdamW(model.parameters(), lr=1e-4)
+            crit   = torch.nn.CrossEntropyLoss()
+            trainer = model.make_trainer(loader, opt, crit)
+            for epoch in range(50):
+                loss = trainer.train_epoch()
+        """
+        try:
+            from trainer import Trainer
+        except ImportError:
+            from SocrateX.trainer import Trainer
+        _device = device or ("cuda" if __import__("torch").cuda.is_available() else "cpu")
+        return Trainer(self, dataloader, optimizer, criterion, device=_device)
+
+    def generate_data(self, source, count=1000, output_dir="silly_train", mode="train"):
+        """
+        Generate a quick synthetic dataset directly from the model object.
+
+        Args:
+            source: URL or file path with words to render
+            count: number of images to generate
+            output_dir: folder to save images + labels.csv
+            mode: 'train' or 'test'
+
+        Example::
+
+            model.generate_data(
+                source="https://raw.githubusercontent.com/.../google-10000-english.txt",
+                count=500,
+                output_dir="my_data",
+                mode="train"
+            )
+        """
+        try:
+            from synthetic import generate_silly_training_set, generate_silly_testing_set
+        except ImportError:
+            from SocrateX.synthetic import generate_silly_training_set, generate_silly_testing_set
+        if mode == "train":
+            return generate_silly_training_set(source=source, count=count, output_dir=output_dir)
+        else:
+            return generate_silly_testing_set(source=source, count=count, output_dir=output_dir)
+
+    def load_data(self, path):
+        """
+        Load a dataset from a CSV / JSON / TXT file.
+        Returns (images, labels).
+
+        Example::
+
+            images, labels = model.load_data("label.csv")
+            dataset = model.make_dataset(images, labels)
+        """
+        try:
+            from dataset import load_dataset
+        except ImportError:
+            from SocrateX.dataset import load_dataset
+        return load_dataset(path)
 
 # Factory functions for models