Update app.py

app.py

@@ -7,14 +7,24 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStream
 MODEL_ID = "HuggingFaceTB/nanowhale-100m"
 
 print(f"Loading model {MODEL_ID} ...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
-    device_map="auto",
-    trust_remote_code=True
-)
-model.eval()
+import torch
+from safetensors.torch import load_file
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+from huggingface_hub import hf_hub_download
+
+# Load model (recommended: manual load for reliability)
+config = AutoConfig.from_pretrained("HuggingFaceTB/nanowhale-100m", trust_remote_code=True)
+model = AutoModelForCausalLM.from_config(config, trust_remote_code=True).float()
+
+# Download and load weights
+weights_path = hf_hub_download("HuggingFaceTB/nanowhale-100m", "model.safetensors")
+state_dict = load_file(weights_path)
+model.load_state_dict(state_dict, strict=True)
+model = model.cuda().eval()
+
+tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/nanowhale-100m")
+
+
 print("Model loaded.")
 
 DEVICE = next(model.parameters()).device