Update app.py

app.py

@@ -4,11 +4,16 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 from fastapi.middleware.cors import CORSMiddleware
 import logging
+from huggingface_hub import HfFolder
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+hf_token = os.getenv('HF_API_TOKEN')
+if not hf_token:
+    raise ValueError("Hugging Face API token is not set. Please set the HF_API_TOKEN environment variable.")
+
 app = FastAPI()
 
 # Enable CORS
@@ -24,8 +29,8 @@ app.add_middleware(
 MODEL_NAME = "aubmindlab/aragpt2-medium"
 
 try:
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token = hf_token)
+    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, token = hf_token)
 except Exception as e:
     logger.error(f"Failed to load model or tokenizer: {str(e)}")
     raise RuntimeError(f"Failed to load model or tokenizer: {str(e)}")