More fixes to HF HOME cache

main.py

@@ -6,22 +6,26 @@ from pydantic import BaseModel
 import os
 
 app = FastAPI()
-os.makedirs("/app/cache", exist_ok=True)
+os.environ["HF_HOME"] = "/tmp/huggingface"
+os.makedirs(os.environ["HF_HOME"], exist_ok=True)
 
 model_path = hf_hub_download(repo_id="MienOlle/sentiment_analysis_api",
                              filename="sentimentAnalysis.pth",
-                             cache_dir="/app/cache"
+                             cache_dir=os.environ["HF_HOME"]
                              )
 modelToken = token.from_pretrained("mdhugol/indonesia-bert-sentiment-classification")
 model = modelSC.from_pretrained("mdhugol/indonesia-bert-sentiment-classification", num_labels=3)
-model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.load_state_dict(torch.load(model_path, map_location=torch.device(device)))
+model.to(device)
 model.eval()
 
 class TextInput(BaseModel):
     text: str
 
 def predict(input):
-    inputs = modelToken(input, return_tensors="pt", padding=True, truncation=True, max_length=512)
+    inputs = modelToken(input, return_tensors="pt", padding=True, truncation=True, max_length=512).to(device)
     
     with torch.no_grad():
         outputs = model(**inputs)