brain gpt

Dockerfile

@@ -1,4 +1,4 @@
-FROM python:3.9
+FROM python:3.10
 
 WORKDIR /code
 

app/model_loader.py

@@ -1,43 +1,13 @@
 
-# import os
-# from transformers import AutoModelForCausalLM, AutoTokenizer
-# from peft import PeftModel
-
-# # Define a safe cache location
-# HF_CACHE = "/tmp/huggingface_cache"
-# os.environ["TRANSFORMERS_CACHE"] = HF_CACHE
-# os.environ["HF_HOME"] = HF_CACHE
-# os.makedirs(HF_CACHE, exist_ok=True)
-
-# def load_model():
-#     hf_token = os.getenv("HF_TOKEN")
-#     if not hf_token:
-#         raise RuntimeError("HF_TOKEN not set.")
-
-#     base_model = AutoModelForCausalLM.from_pretrained(
-#         "meta-llama/Llama-2-7b-chat-hf",
-#         use_auth_token=hf_token,
-#         cache_dir="/tmp/hf_cache",
-#         torch_dtype="auto",
-#         device_map="auto"
-#     )
-#     model = PeftModel.from_pretrained(
-#         base_model,
-#         "BrainGPT/BrainGPT-7B-v0.1",
-#         use_auth_token=hf_token,
-#         cache_dir="/tmp/hf_cache"
-#     )
-#     tokenizer = AutoTokenizer.from_pretrained(
-#         "meta-llama/Llama-2-7b-chat-hf",
-#         use_auth_token=hf_token,
-#         cache_dir="/tmp/hf_cache"
-#     )
-#     return model, tokenizer
-
 import os
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
 
 def load_model():
+    hf_token = os.getenv("HF_TOKEN")
+    if not hf_token:
+        raise RuntimeError("HF_TOKEN not set.")
+    
     # Use a user-writable cache directory (important for Docker non-root)
     HF_CACHE = os.path.expanduser("~/.cache/huggingface")
     os.makedirs(HF_CACHE, exist_ok=True)
@@ -45,16 +15,47 @@ def load_model():
     os.environ["TRANSFORMERS_CACHE"] = HF_CACHE
     os.environ["HF_HOME"] = HF_CACHE
 
-    model_name = "gpt2"
-
+    base_model = AutoModelForCausalLM.from_pretrained(
+        "meta-llama/Llama-2-7b-chat-hf",
+        use_auth_token=hf_token,
+        cache_dir="/tmp/hf_cache",
+        torch_dtype="auto",
+        device_map="auto"
+    )
+    model = PeftModel.from_pretrained(
+        base_model,
+        "BrainGPT/BrainGPT-7B-v0.1",
+        use_auth_token=hf_token,
+        cache_dir="/tmp/hf_cache"
+    )
     tokenizer = AutoTokenizer.from_pretrained(
-        model_name,
-        cache_dir=HF_CACHE
+        "meta-llama/Llama-2-7b-chat-hf",
+        use_auth_token=hf_token,
+        cache_dir="/tmp/hf_cache"
     )
+    return model, tokenizer
 
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        cache_dir=HF_CACHE
-    )
+# import os
+# from transformers import AutoModelForCausalLM, AutoTokenizer
 
-    return model, tokenizer
+# def load_model():
+#     # Use a user-writable cache directory (important for Docker non-root)
+#     HF_CACHE = os.path.expanduser("~/.cache/huggingface")
+#     os.makedirs(HF_CACHE, exist_ok=True)
+
+#     os.environ["TRANSFORMERS_CACHE"] = HF_CACHE
+#     os.environ["HF_HOME"] = HF_CACHE
+
+#     model_name = "gpt2"
+
+#     tokenizer = AutoTokenizer.from_pretrained(
+#         model_name,
+#         cache_dir=HF_CACHE
+#     )
+
+#     model = AutoModelForCausalLM.from_pretrained(
+#         model_name,
+#         cache_dir=HF_CACHE
+#     )
+
+#     return model, tokenizer