Create model_loader.py

model_loader.py

@@ -0,0 +1,28 @@
+import os
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+
+HF_KEY = os.getenv("HF_Key")
+MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"  # bạn có thể đổi model nhỏ hơn nếu chậm
+
+def load_model():
+    print("Đang load model:", MODEL_NAME)
+
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=HF_KEY)
+
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        torch_dtype=torch.float32,       # CPU nên dùng float32, nếu GPU có thể dùng float16
+        low_cpu_mem_usage=True,
+        device_map="cpu",
+        token=HF_KEY
+    )
+
+    pipe = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+    )
+
+    print("Model loaded thành công")
+    return pipe