Update app.py

app.py

@@ -204,18 +204,12 @@ def load_model():
         print("Adding special tokens...")
         tokenizer.add_special_tokens(special_tokens)
         
-        print("Loading model configuration...")
-        config = SmolLM2Config()
-        
-        print("Initializing model...")
-        model = SmolLM2ForCausalLM(config)
-        
-        print("Loading model weights...")
-        state_dict = torch.load(
-            os.path.join(model_id, "pytorch_model.bin"),
-            map_location="cpu"
+        print("Loading model from Hugging Face Hub...")
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            trust_remote_code=True
         )
-        model.load_state_dict(state_dict)
         
         # Move model to device manually
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")