Update main.py

main.py

@@ -19,7 +19,11 @@ import tempfile
 os.environ["HF_HOME"] = "/tmp/huggingface"
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface/transformers"
 os.environ["HF_HUB_CACHE"] = "/tmp/huggingface/hub"
-os.makedirs(os.environ["HF_HOME"], exist_ok=True)
+os.environ["HF_DATASETS_CACHE"] = "/tmp/huggingface/datasets"
+os.environ["XDG_CACHE_HOME"] = "/tmp/huggingface"
+
+for path in os.environ.values():
+    os.makedirs(path, exist_ok=True)
 
 # Silence all transformers and huggingface logging
 logging.getLogger("transformers").setLevel(logging.ERROR)
@@ -30,19 +34,22 @@ app = Flask(__name__)
 CORS(app)
 
 # ========== Load Whisper Model (quantized) ==========
-def load_whisper_model(model_size="small"):
+def load_whisper_model(model_size="small", save_dir="/tmp/models_cache/whisper"):
+    os.makedirs(save_dir, exist_ok=True)
     model_name = f"openai/whisper-{model_size}"
-    processor = WhisperProcessor.from_pretrained(model_name)
-    model = WhisperForConditionalGeneration.from_pretrained(model_name)
+    processor = WhisperProcessor.from_pretrained(model_name, cache_dir=save_dir)
+    model = WhisperForConditionalGeneration.from_pretrained(model_name, cache_dir=save_dir)
     model = quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)
     model.to("cuda" if torch.cuda.is_available() else "cpu")
     return processor, model
 
+
 # ========== Load Grammar Correction Model (quantized) ==========
-def load_grammar_model():
+def load_grammar_model(save_dir="/tmp/models_cache/grammar_corrector"):
+    os.makedirs(save_dir, exist_ok=True)
     model_name = "prithivida/grammar_error_correcter_v1"
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+    tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=save_dir)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name, cache_dir=save_dir)
     model = quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)
     grammar_pipeline = pipeline(
         "text2text-generation",