added inference code to README

README.md

@@ -53,6 +53,156 @@ Here's the training loss progression:
 
 ![Training Loss Curve](plots/pretrain_loss_20250303.png)
 
+### Inference
+
+```pyrhon
+from huggingface_hub import snapshot_download
+snapshot_download(repo_id="PursuitOfDataScience/Argonne-1.0")
+```
+
+Set up `minimal_chat.py` as follows:
+
+```python
+import os
+import sys
+import torch
+import json
+import argparse
+import time
+from transformers import AutoTokenizer
+
+def main():
+    parser = argparse.ArgumentParser(description="Minimal Argonne chat")
+    parser.add_argument("--model_dir", required=True, help="Directory containing model files")
+    parser.add_argument("--mp_dir", required=True, help="Directory containing mp_pretrain.py")
+    args = parser.parse_args()
+    
+    # Print all input arguments
+    print(f"Model directory: {args.model_dir}")
+    print(f"mp_pretrain directory: {args.mp_dir}")
+    
+    # Check that directories exist
+    if not os.path.exists(args.model_dir):
+        print(f"Error: Model directory {args.model_dir} does not exist")
+        sys.exit(1)
+    if not os.path.exists(args.mp_dir):
+        print(f"Error: mp_pretrain directory {args.mp_dir} does not exist")
+        sys.exit(1)
+        
+    # Check for required files
+    required_files = [
+        os.path.join(args.model_dir, "config.json"),
+        os.path.join(args.model_dir, "tokenizer.json")
+    ]
+    
+    for file_path in required_files:
+        if not os.path.exists(file_path):
+            print(f"Error: Required file {file_path} does not exist")
+            sys.exit(1)
+    
+    # Check for either pytorch_model.bin or model.safetensors
+    weights_file = None
+    if os.path.exists(os.path.join(args.model_dir, "pytorch_model.bin")):
+        weights_file = os.path.join(args.model_dir, "pytorch_model.bin")
+        print(f"Found PyTorch weights at {weights_file}")
+    elif os.path.exists(os.path.join(args.model_dir, "model.safetensors")):
+        weights_file = os.path.join(args.model_dir, "model.safetensors")
+        print(f"Found safetensors weights at {weights_file}")
+    else:
+        print(f"Error: No model weights found in {args.model_dir}")
+        sys.exit(1)
+        
+    # Add mp_pretrain directory to Python path
+    sys.path.insert(0, args.mp_dir)
+    
+    # Import required modules
+    try:
+        print("Importing modules from mp_pretrain...")
+        from mp_pretrain import ArgonneModelParallel, ArgonneConfig, load_bpe_tokenizer
+        print("Import successful")
+    except ImportError as e:
+        print(f"Error importing modules from mp_pretrain.py: {e}")
+        sys.exit(1)
+        
+    # Load the config
+    print("Loading model config...")
+    with open(os.path.join(args.model_dir, "config.json"), 'r') as f:
+        config_dict = json.load(f)
+    config = ArgonneConfig(**config_dict)
+    print("Config loaded")
+    
+    # Load the tokenizer
+    print("Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(args.model_dir)
+    print("Tokenizer loaded")
+    
+    # Create the model
+    print("Creating model...")
+    model = ArgonneModelParallel(config)
+    print("Model created")
+    
+    # Load weights
+    print(f"Loading weights from {weights_file}...")
+    if weights_file.endswith(".bin"):
+        # Load PyTorch weights
+        state_dict = torch.load(weights_file, map_location="cpu")
+    else:
+        # Load safetensors weights
+        from safetensors.torch import load_file
+        state_dict = load_file(weights_file)
+    
+    # Load state dict
+    print("Applying weights to model...")
+    model.load_state_dict(state_dict, strict=False)
+    print("Weights loaded")
+    
+    # Move to GPU if available
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Moving model to {device}...")
+    model = model.to(device)
+    
+    # Set devices attribute needed for generate
+    model.devices = [device]
+    
+    print("Model ready for chat!")
+    
+    # Simple chat loop
+    print("\n" + "="*50)
+    print("Argonne Model Chat - Type 'exit' to quit")
+    print("="*50 + "\n")
+    
+    while True:
+        user_input = input("You: ").strip()
+        if user_input.lower() in ["exit", "quit"]:
+            print("Goodbye!")
+            break
+            
+        # Generate response
+
+        # Encode input
+        input_ids = tokenizer.encode(user_input, return_tensors="pt").to(device)
+        
+        # Generate
+        with torch.no_grad():
+            output_ids = model.generate(
+                input_ids,
+                max_new_tokens=50,
+                temperature=0.7,
+                top_k=50)[0]
+            
+        # Decode output
+        response = tokenizer.decode(output_ids, skip_special_tokens=True)
+        print(f"Model: {response}")
+
+
+if __name__ == "__main__":
+    main()
+
+```
+
+```
+python minimal_chat.py --model_dir /path/to/model --mp_dir /path/to/mp_pretrain.py
+```
 
 ### 📝 Example Outputs
 Below are generated examples illustrating Argonne-1.0's style and capability when prompted: