From 63c341049254ed6c64136d17cbad5299ee8729ee Mon Sep 17 00:00:00 2001
From: teleprint-me <77757836+teleprint-me@users.noreply.github.com>
Date: Sat, 25 May 2024 04:15:39 -0400
Subject: [PATCH] refactor: Add support for model file types

---
 gguf-py/gguf/constants.py | 24 +++++++++++++++---------
 1 file changed, 15 insertions(+), 9 deletions(-)

diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py
index 48f055a15..382f0fda2 100644
--- a/gguf-py/gguf/constants.py
+++ b/gguf-py/gguf/constants.py
@@ -969,18 +969,24 @@ VOCAB_TYPE_NAMES: dict[VocabType, str] = {
 # Model File Types
 #
 class ModelFileType(IntEnum):
-    UNK = auto()  # Unsupported file type
-    SFT = auto()  # SafeTensor file type
-    PTH = auto()  # PyTorch file type
-    BIN = auto()  # Pickled file type
-    PT  = auto()  # PyTorch file type
+    UNK         = auto()
+    PT          = auto()
+    PTH         = auto()
+    BIN         = auto()
+    SAFETENSORS = auto()
+    JSON        = auto()
+    MODEL       = auto()
+    GGUF        = auto()
 
 
 MODEL_FILE_TYPE_NAMES: dict[ModelFileType, str] = {
-    ModelFileType.PT:  "pt",
-    ModelFileType.PTH: "pth",
-    ModelFileType.BIN: "bin",
-    ModelFileType.SFT: "safetensors",
+    ModelFileType.PT:          ".pt",
+    ModelFileType.PTH:         ".pth",
+    ModelFileType.BIN:         ".bin",
+    ModelFileType.SAFETENSORS: ".safetensors",
+    ModelFileType.JSON:        ".json",
+    ModelFileType.MODEL:       ".model",
+    ModelFileType.GGUF:        ".gguf",
 }
 
 #