Update main.py

main.py

@@ -7,18 +7,25 @@ from typing import Optional, Any
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextStreamer, GenerationConfig
 import os
+from huggingface_hub import login
+
+# Check whether we are executing inside a Hugging Face Space
+SPACE_NAME = os.getenv("SPACE_NAME", default=None)
+if SPACE_NAME is not None:
+    print(f"Running inside {SPACE_NAME} Space.")
+    try:
+        # Try to auto-login using the Space's environment variables
+        login(automatically=True)
+    except Exception as e:
+        print(f"Failed to auto-login ({str(e)}). Manually check the HF_ACCESS_TOKEN environment variable.")
+        sys.exit(1)
 
 try:
-    # Attempt to fetch the token from the environment variable set in Spaces
     HUGGINGFACE_TOKEN = os.environ['HF_ACCESS_TOKEN']
 except KeyError:
     print('The environment variable "HF_ACCESS_TOKEN" is not found. Please configure it correctly in your Space.')
     sys.exit(1)
 
-# Authentication
-from huggingface_hub import interpreter_login
-interpreter_login()
-
 # Packages and model loading
 import torch
 base_model_id = "152334H/miqu-1-70b-sf"