Updated codet5_model.py

codet5_model.py

@@ -34,8 +34,11 @@
 #     output = model.generate(**inputs, max_length=256)
 #     return tokenizer.decode(output[0], skip_special_tokens=True)
 
+import os
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
+MODEL_DIR = "/data/codet5-model"
+
 # Global variables but not initialized
 tokenizer = None
 model = None
@@ -43,9 +46,14 @@ model = None
 def load_model():
     global tokenizer, model
     if tokenizer is None or model is None:
-        print("Loading CodeT5 model...")
-        tokenizer = AutoTokenizer.from_pretrained("Salesforce/codet5p-770m-py")
-        model = AutoModelForSeq2SeqLM.from_pretrained("Salesforce/codet5p-770m-py")
+        if not os.path.exists(MODEL_DIR):
+            print("Downloading CodeT5 model to persistent /data directory...")
+            tokenizer = AutoTokenizer.from_pretrained("Salesforce/codet5p-770m-py", cache_dir=MODEL_DIR)
+            model = AutoModelForSeq2SeqLM.from_pretrained("Salesforce/codet5p-770m-py", cache_dir=MODEL_DIR)
+        else:
+            print("Loading model from /data directory...")    
+            tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
+            model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_DIR)
         print("Model loaded.")
 
 def generate_code(instruction: str) -> str:

whisper_model.py

@@ -1,12 +1,13 @@
 import os
+import whisper
+import logging
 
 # Add FFmpeg directory to PATH
-os.environ["PATH"] += os.pathsep + r"C:\ffmpeg\bin"
+# os.environ["PATH"] += os.pathsep + r"C:\ffmpeg\bin" 
 
-import whisper
-import logging
+WHISPER_MODEL_DIR = "/data/whisper-small"
 
-model = whisper.load_model("small")
+model = whisper.load_model("small", download_root=WHISPER_MODEL_DIR)
 
 # def transcribe_audio(file_path: str) -> str:
 #   result = model.transcribe(file_path)