Upload model

config.json

@@ -1,4 +1,11 @@
 {
+  "architectures": [
+    "SafeSearchModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "model.SafeSearchConfig",
+    "AutoModelForImageClassification": "model.SafeSearchModel"
+  },
   "classes": {
     "0": "nsfw_gore",
     "1": "nsfw_suggestive",
@@ -6,6 +13,7 @@
   },
   "classifier": "default",
   "crop_pct": 0.875,
+  "device": "cpu",
   "first_conv": "conv2d",
   "has_aux": false,
   "input_channels": 3,
@@ -34,5 +42,6 @@
     0.5,
     0.5
   ],
-  "transformers_version": "4.26.1"
-}
+  "torch_dtype": "float32",
+  "transformers_version": "4.25.1"
+}

model.py

@@ -0,0 +1,62 @@
+classes = { '0': 'nsfw_gore', '1': 'nsfw_suggestive', '2': 'safe' }
+model_path = "safesearch_mini_v2.bin"
+
+from transformers import PretrainedConfig, PreTrainedModel
+
+class SafeSearchConfig(PretrainedConfig):
+	model_type = "safesearch_mini_v2"
+	def __init__(self,
+				 model_name: str = "safesearch_mini_v2",
+				 input_channels: int = 3,
+				 num_classes: int = 3,
+				 input_size: list[int] = [3, 299, 299],
+				 pool_size: list[int] = [8, 8],
+				 crop_pct: float = 0.875,
+				 interpolation: str = "bicubic",
+				 mean: list[float] = [0.5, 0.5, 0.5],
+				 std: list[float] = [0.5, 0.5, 0.5],
+				 first_conv: str = "conv2d",
+				 classifier: str = "default",
+				 has_aux: bool = False,
+				 label_offset: int = 0,
+				 classes: list[str] = classes,
+				 output_channels: int = 1536,
+				 device: str = "cpu",
+				 **kwargs):
+		self.model_name = model_name
+		self.input_channels = input_channels
+		self.num_classes = num_classes
+		self.input_size = input_size
+		self.pool_size = pool_size
+		self.crop_pct = crop_pct
+		self.interpolation = interpolation
+		self.mean = mean
+		self.std = std
+		self.first_conv = first_conv
+		self.classifier = classifier
+		self.has_aux = has_aux
+		self.label_offset = label_offset
+		self.classes = classes
+		self.output_channels = output_channels
+		self.device = device
+		super().__init__(**kwargs)
+
+"""
+safesearch_config = SafeSearchConfig()
+safesearch_config.save_pretrained("safesearch_config")
+"""
+
+import torch, os, timm
+
+class SafeSearchModel(PreTrainedModel):
+	config_class = SafeSearchConfig
+	def __init__(self, config: SafeSearchConfig):
+		super().__init__(config)
+		if not os.path.exists(model_path):
+			from urllib.request import urlretrieve
+			urlretrieve(f"https://huggingface.co/FredZhang7/google-safesearch-mini-v2/resolve/main/pytorch_model.bin", model_path)
+		self.model = timm.create_model("inception_resnet_v2", pretrained=False, num_classes=3)
+		self.model.load_state_dict(torch.load(model_path, map_location=torch.device(config.device)))
+
+	def forward(self, input_ids: torch.Tensor):
+		return self.model(input_ids)

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:73b58de7e47cdb1d0bbca9ff1978da93415d7cbb3930465488d9d745ddc744f2
-size 217948223
+oid sha256:55829a7fa31ff8f643b9c7158c449573c52ee3f4ab52b87a669b920a503348ee
+size 217901159