Upload 4 files

aglm/config.json

@@ -0,0 +1,12 @@
+{
+  "model_name": "AgLM",
+  "block_size": 128,
+  "n_embd": 128,
+  "n_head": 4,
+  "n_layer": 4,
+  "vocab_size": 8000,
+  "batch_size": 8,
+  "grad_accum": 4,
+  "max_epochs": 3,
+  "end_token_id": 4
+}

aglm/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b3ce58b730f490c19b272866d043e0ff9eb64abb648a392d0a2eb35314211d5a
+size 12512854

gpt_chat.py

@@ -0,0 +1,141 @@
+import torch
+import json
+from torch.nn import functional as F
+from tokenizers import Tokenizer
+from pathlib import Path
+import os
+
+# LightweightGPT Model
+class LightweightGPT(torch.nn.Module):
+    """Compact GPT-like model with causal masking and positional encoding"""
+    def __init__(self, vocab_size, block_size, n_embd, n_head, n_layer):
+        super().__init__()
+        self.block_size = block_size
+        self.token_embedding = torch.nn.Embedding(vocab_size, n_embd)
+        self.position_embedding = torch.nn.Embedding(block_size, n_embd)
+        
+        self.blocks = torch.nn.ModuleList([
+            torch.nn.TransformerDecoderLayer(
+                d_model=n_embd,
+                nhead=n_head,
+                dim_feedforward=4 * n_embd,
+                dropout=0.1,
+                activation='gelu',
+                batch_first=True,
+                norm_first=True
+            )
+            for _ in range(n_layer)
+        ])
+        self.ln_f = torch.nn.LayerNorm(n_embd)
+        self.lm_head = torch.nn.Linear(n_embd, vocab_size, bias=False)
+
+    def forward(self, idx, targets=None):
+        B, T = idx.shape
+        device = idx.device
+        causal_mask = torch.triu(torch.ones(T, T, device=device, dtype=torch.bool), diagonal=1)
+
+        token_emb = self.token_embedding(idx)
+        pos = torch.arange(0, T, dtype=torch.long, device=device)
+        pos_emb = self.position_embedding(pos)
+        
+        x = token_emb + pos_emb
+
+        for block in self.blocks:
+            x = block(x, x, tgt_mask=causal_mask)
+
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(
+                logits.view(-1, logits.size(-1)),
+                targets.view(-1),
+                ignore_index=-1
+            )
+        return logits, loss
+
+    def generate(self, idx, max_new_tokens, temperature=0.8, top_k=50, stop_token=None):
+        """Generate text with context handling and positional encoding"""
+        for _ in range(max_new_tokens):
+            idx_cond = idx[:, -self.block_size:]
+            logits, _ = self(idx_cond)
+            logits = logits[:, -1, :]
+            logits = logits / temperature
+
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float('Inf')
+
+            probs = F.softmax(logits, dim=-1)
+            idx_next = torch.multinomial(probs, num_samples=1)
+
+            if stop_token is not None and idx_next.item() == stop_token:
+                break
+
+            idx = torch.cat((idx, idx_next), dim=1)
+            
+        return idx
+
+# Chat Interface
+class ChatInterface:
+    def __init__(self, model_dir="/"):
+        self.model_dir = Path(model_dir)
+        self.device = "mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu"
+        self.load_model()
+        
+    def load_model(self):
+        with open(self.model_dir / "config.json", "r") as f:
+            self.config = json.load(f)
+        
+        self.tokenizer = Tokenizer.from_file(str(self.model_dir / "tokenizer.json"))
+        self.end_token_id = self.config.get("end_token_id")
+        
+        self.model = LightweightGPT(
+            vocab_size=self.config["vocab_size"],
+            block_size=self.config["block_size"],
+            n_embd=self.config["n_embd"],
+            n_head=self.config["n_head"],
+            n_layer=self.config["n_layer"]
+        ).to(self.device)
+        
+        self.model.load_state_dict(torch.load(self.model_dir / "model.pt", map_location=self.device))
+        self.model.eval()
+        print("✅ Model loaded successfully!")
+    
+    def chat(self):
+        print("\n===== AI Assistant Ready =====")
+        print("Type 'quit' or 'exit' to end the chat.\n")
+        
+        while True:
+            user_input = input("user: ")
+            if user_input.lower() in ["quit", "exit"]:
+                break
+                
+            prompt = f"user: {user_input}\nai:"
+            input_ids = self.tokenizer.encode(prompt).ids
+            input_tensor = torch.tensor([input_ids], dtype=torch.long, device=self.device)
+            
+            with torch.no_grad():
+                output_ids = self.model.generate(
+                    input_tensor,
+                    max_new_tokens=150,
+                    temperature=0.7,
+                    top_k=40,
+                    stop_token=self.end_token_id
+                )
+            
+            response_ids = output_ids[0, len(input_ids):].tolist()
+            response = self.tokenizer.decode(response_ids)
+            response = response.replace("<|endoftext|>", "").strip()
+            
+            print(f"ai: {response}")
+
+# Main execution
+if __name__ == "__main__":
+    model_folder = "aglm"
+    if os.path.exists(model_folder) and os.path.exists(os.path.join(model_folder, "model.pt")):
+        chat_bot = ChatInterface(model_dir=model_folder)
+        chat_bot.chat()
+    else:
+        print(f"\nERROR: Model directory '{model_folder}' not found. Please ensure the model is trained and the directory contains 'model.pt' and 'config.json'.")