Create model_loader.py

model_loader.py

@@ -0,0 +1,26 @@
+# model_loader.py
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from safetensors.torch import load_file
+import torch
+from config import MODEL_NAME, DEVICE
+
+def load_model():
+    """
+    Load model using safetensors if available, else fallback to normal HF model.
+    """
+    try:
+        # Attempt to load safetensors
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
+            device_map="auto",
+            torch_dtype=torch.float16,
+            use_safetensors=True
+        )
+        print("[INFO] Model loaded using safetensors.")
+    except Exception as e:
+        print(f"[WARN] Could not use safetensors. Loading normal model. {e}")
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+        model = AutoModelForCausalLM.from_pretrained(MODEL_NAME).to(DEVICE)
+    
+    return tokenizer, model