Upload folder using huggingface_hub

inference.py

@@ -1,35 +1,92 @@
 import torch
 from torchvision import transforms
-from pathlib import Path
-import json
-import os
-import sys
 from model import MNISTModel
-from inference_util import Inferencer
 
 
 class InferenceWrapper:
-    def __init__(self, model_path: str, input_dir: str = 'input_data', output_dir: str = 'output_data'):
+    def __init__(self, model_path: str):
+        """
+        Initialize the inference wrapper with a model path.
+        
+        Args:
+            model_path (str): Path to the model weights file
+        """
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         self.model_path = model_path
-        self.inferencer = Inferencer(input_dir, output_dir)
-        # Override the model with our specified model path
-        self.inferencer.model, _ = self.inferencer._load_model(model_path)
+        self.model = self._load_model()
+        self.transform = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize((0.1307,), (0.3081,))
+        ])
+        
+    def _load_model(self):
+        """Load and return the model."""
+        model = MNISTModel().to(self.device)
+        model.load_state_dict(
+            torch.load(self.model_path, map_location=self.device, weights_only=True)
+        )
+        model.eval()
+        return model
+
+    def predict_tensor(self, input_tensor: torch.Tensor):
+        """
+        Run inference on a single input tensor.
+        
+        Args:
+            input_tensor (torch.Tensor): Input tensor of shape [1, 28, 28] or [N, 1, 28, 28]
+            
+        Returns:
+            tuple: (prediction, confidence)
+        """
+        with torch.no_grad():
+            if input_tensor.dim() == 3:
+                input_tensor = input_tensor.unsqueeze(0)
+            
+            input_tensor = input_tensor.to(self.device)
+            output = self.model(input_tensor)
+            probs = torch.softmax(output, dim=1)
+            prediction = output.argmax(1).item()
+            confidence = probs[0][prediction].item()
+            return prediction, confidence
+
+    def predict_batch(self, input_tensors: torch.Tensor):
+        """
+        Run inference on a batch of input tensors.
+        
+        Args:
+            input_tensors (torch.Tensor): Batch of input tensors of shape [N, 1, 28, 28]
+            
+        Returns:
+            tuple: (predictions, confidences)
+        """
+        with torch.no_grad():
+            input_tensors = input_tensors.to(self.device)
+            output = self.model(input_tensors)
+            probs = torch.softmax(output, dim=1)
+            predictions = output.argmax(1)
+            confidences = torch.gather(probs, 1, predictions.unsqueeze(1)).squeeze(1)
+            return predictions.cpu().numpy(), confidences.cpu().numpy()
 
-    def run_inference(self):
-        """Run inference using the specified model"""
-        return self.inferencer.process_input()
 
 def main():
     import argparse
     parser = argparse.ArgumentParser()
     parser.add_argument('--model-path', required=True, help='Path to the model weights')
-    parser.add_argument('--input-dir', default='input_data')
-    parser.add_argument('--output-dir', default='output_data')
     args = parser.parse_args()
     
-    wrapper = InferenceWrapper(args.model_path, args.input_dir, args.output_dir)
-    results = wrapper.run_inference()
-    print(f"Processed {len(results)} inputs using model: {args.model_path}")
+    # Example usage
+    wrapper = InferenceWrapper(args.model_path)
+    
+    # Example single inference
+    test_input = torch.randn(1, 28, 28)
+    prediction, confidence = wrapper.predict_tensor(test_input)
+    print(f"Single prediction: {prediction}, confidence: {confidence:.4f}")
+    
+    # Example batch inference
+    batch_input = torch.randn(4, 1, 28, 28)
+    predictions, confidences = wrapper.predict_batch(batch_input)
+    print(f"Batch predictions: {predictions}")
+    print(f"Batch confidences: {confidences}")
 
 if __name__ == "__main__":
-    main()
+    main()

mnist/.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

mnist/best_model.pth

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

mnist/config.json

@@ -0,0 +1,9 @@
+{
+    "model_name": "mnist_classifier",
+    "input_size": [1, 28, 28],
+    "hidden_size": 128,
+    "num_classes": 10,
+    "dropout": 0.5,
+    "mean": 0.1307,
+    "std": 0.3081
+}

mnist/inference.py

@@ -0,0 +1,92 @@
+import torch
+from torchvision import transforms
+from model import MNISTModel
+
+
+class InferenceWrapper:
+    def __init__(self, model_path: str):
+        """
+        Initialize the inference wrapper with a model path.
+        
+        Args:
+            model_path (str): Path to the model weights file
+        """
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.model_path = model_path
+        self.model = self._load_model()
+        self.transform = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize((0.1307,), (0.3081,))
+        ])
+        
+    def _load_model(self):
+        """Load and return the model."""
+        model = MNISTModel().to(self.device)
+        model.load_state_dict(
+            torch.load(self.model_path, map_location=self.device, weights_only=True)
+        )
+        model.eval()
+        return model
+
+    def predict_tensor(self, input_tensor: torch.Tensor):
+        """
+        Run inference on a single input tensor.
+        
+        Args:
+            input_tensor (torch.Tensor): Input tensor of shape [1, 28, 28] or [N, 1, 28, 28]
+            
+        Returns:
+            tuple: (prediction, confidence)
+        """
+        with torch.no_grad():
+            if input_tensor.dim() == 3:
+                input_tensor = input_tensor.unsqueeze(0)
+            
+            input_tensor = input_tensor.to(self.device)
+            output = self.model(input_tensor)
+            probs = torch.softmax(output, dim=1)
+            prediction = output.argmax(1).item()
+            confidence = probs[0][prediction].item()
+            return prediction, confidence
+
+    def predict_batch(self, input_tensors: torch.Tensor):
+        """
+        Run inference on a batch of input tensors.
+        
+        Args:
+            input_tensors (torch.Tensor): Batch of input tensors of shape [N, 1, 28, 28]
+            
+        Returns:
+            tuple: (predictions, confidences)
+        """
+        with torch.no_grad():
+            input_tensors = input_tensors.to(self.device)
+            output = self.model(input_tensors)
+            probs = torch.softmax(output, dim=1)
+            predictions = output.argmax(1)
+            confidences = torch.gather(probs, 1, predictions.unsqueeze(1)).squeeze(1)
+            return predictions.cpu().numpy(), confidences.cpu().numpy()
+
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--model-path', required=True, help='Path to the model weights')
+    args = parser.parse_args()
+    
+    # Example usage
+    wrapper = InferenceWrapper(args.model_path)
+    
+    # Example single inference
+    test_input = torch.randn(1, 28, 28)
+    prediction, confidence = wrapper.predict_tensor(test_input)
+    print(f"Single prediction: {prediction}, confidence: {confidence:.4f}")
+    
+    # Example batch inference
+    batch_input = torch.randn(4, 1, 28, 28)
+    predictions, confidences = wrapper.predict_batch(batch_input)
+    print(f"Batch predictions: {predictions}")
+    print(f"Batch confidences: {confidences}")
+
+if __name__ == "__main__":
+    main()

mnist/inference_util.py

@@ -0,0 +1,87 @@
+# inference.py
+import torch
+from torchvision import transforms, datasets
+from PIL import Image
+import json
+from pathlib import Path
+from model import MNISTModel
+import os
+import sys
+
+class Inferencer:
+    def __init__(self, input_dir: str = 'input_data', output_dir: str = 'output_data'):
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.model, _ = self._load_model()
+        self.input_dir = Path(input_dir)
+        self.output_dir = Path(output_dir)
+        self.transform = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize((0.1307,), (0.3081,))
+        ])
+        
+    def _load_model(self, model_path='best_model.pth'):
+        """Load the trained model."""
+        model = MNISTModel().to(self.device)
+        model.load_state_dict(
+            torch.load(model_path, map_location=self.device, weights_only=True)
+        )
+        model.eval()
+        return model, self.device
+    
+    def predict(self, input_tensor: torch.Tensor):
+        """Make prediction on the input tensor."""
+        with torch.no_grad():
+            if input_tensor.dim() == 3:
+                input_tensor = input_tensor.unsqueeze(0)
+            
+            input_tensor = input_tensor.to(self.device)
+            output = self.model(input_tensor)
+            probs = torch.softmax(output, dim=1)
+            prediction = output.argmax(1).item()
+            confidence = probs[0][prediction].item()
+            return prediction, confidence
+
+    def process_input(self):
+        """Process all images in input directory."""
+        # Create output directory if it doesn't exist
+        os.makedirs(self.output_dir, exist_ok=True)
+        
+        results = []
+        # Process each file in input directory
+        for file_path in sorted(self.input_dir.glob('*.pt')):  # For tensor files
+            try:
+                # Load tensor
+                input_tensor = torch.load(file_path)
+                
+                # Get prediction
+                prediction, confidence = self.predict(input_tensor)
+                
+                results.append({
+                    "filename": file_path.name,
+                    "prediction": prediction,
+                    "confidence": confidence
+                })
+                
+            except Exception as e:
+                print(f"Error processing {file_path}: {str(e)}", file=sys.stderr)
+        
+        # Save results
+        with open(self.output_dir / 'results.json', 'w') as f:
+            json.dump(results, f, indent=2)
+        
+        return results
+
+def main():
+    # Accept input/output directories as arguments
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--input-dir', default='input_data')
+    parser.add_argument('--output-dir', default='output_data')
+    args = parser.parse_args()
+    
+    inferencer = Inferencer(args.input_dir, args.output_dir)
+    results = inferencer.process_input()
+    print(f"Processed {len(results)} inputs")
+
+if __name__ == "__main__":
+    main()

mnist/model.py

@@ -0,0 +1,16 @@
+import torch
+import torch.nn as nn
+
+class MNISTModel(nn.Module):
+    def __init__(self):
+        super(MNISTModel, self).__init__()
+        self.fc1 = nn.Linear(28 * 28, 128)  # MNIST images are 28x28
+        self.fc2 = nn.Linear(128, 10)
+        self.dropout = nn.Dropout(0.5)
+
+    def forward(self, x):
+        x = x.view(-1, 28 * 28)  # Flatten the input
+        x = torch.relu(self.fc1(x))
+        x = self.dropout(x)
+        x = self.fc2(x)
+        return x

mnist/random2.txt

@@ -0,0 +1 @@
+[!] Failed to build directory tree: No such file or directory (os error 2)