Updated codet5_model.py

codet5_model.py

@@ -34,33 +34,67 @@
 #     output = model.generate(**inputs, max_length=256)
 #     return tokenizer.decode(output[0], skip_special_tokens=True)
 
+# import os
+# from huggingface_hub import snapshot_download
+# from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+# MODEL_DIR = snapshot_download(repo_id="Salesforce/codet5p-770m-py", cache_dir="./data")
+
+# os.makedirs(MODEL_DIR, exist_ok=True)
+# # Global variables but not initialized
+# tokenizer = None
+# model = None
+
+# def load_model():
+#     global tokenizer, model
+#     if tokenizer is None or model is None:
+#         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:
+#     load_model()  # Ensure model is loaded before use
+#     prompt = f"# Instruction: {instruction}\n# Python Code:\n"
+
+#     inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
+#     output = model.generate(**inputs, max_length=256)
+#     return tokenizer.decode(output[0], skip_special_tokens=True)
+
+
 import os
+from huggingface_hub import snapshot_download
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-MODEL_DIR = "./data/codet5-model"
-
+# Constants
+HF_MODEL_ID = "Salesforce/codet5p-770m-py"
+MODEL_DIR = "./data/codet5_model"  # More explicit and organized
 os.makedirs(MODEL_DIR, exist_ok=True)
-# Global variables but not initialized
+
+# Snapshot download (will cache in MODEL_DIR)
+snapshot_download(repo_id=HF_MODEL_ID, cache_dir=MODEL_DIR, local_dir=MODEL_DIR, local_dir_use_symlinks=False)
+
+# Globals
 tokenizer = None
 model = None
 
 def load_model():
     global tokenizer, model
     if tokenizer is None or model is None:
-        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.")
+        print("Loading CodeT5 model from persistent directory...")
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
+        model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_DIR)
+        print("Model loaded successfully.")
 
 def generate_code(instruction: str) -> str:
-    load_model()  # Ensure model is loaded before use
+    load_model()
     prompt = f"# Instruction: {instruction}\n# Python Code:\n"
 
     inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
-    output = model.generate(**inputs, max_length=256)
-    return tokenizer.decode(output[0], skip_special_tokens=True)
+    outputs = model.generate(**inputs, max_length=256)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)