scorevision: push artifact

download_model.py

@@ -0,0 +1,75 @@
+from huggingface_hub import snapshot_download
+import os
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+def download_model():
+    """
+    Tải mô hình Mistral-7B-Instruct từ Hugging Face Hub.
+    Trả về đường dẫn đến thư mục chứa mô hình đã tải.
+    """
+    print("Bắt đầu tải mô hình Mistral-7B-Instruct...")
+    
+    # Đặt biến môi trường HUGGING_FACE_HUB_TOKEN nếu bạn có token
+    hf_token = os.getenv("HUGGINGFACE_API_KEY")
+    
+    # Tải mô hình từ Hugging Face Hub
+    model_path = snapshot_download(
+        repo_id="mistralai/Mistral-7B-Instruct-v0.2",
+        token=hf_token,
+        revision="main"
+    )
+    
+    print(f"Đã tải mô hình từ Hugging Face Hub vào: {model_path}")
+    
+    # Kiểm tra mô hình đã tải
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_path)
+        print("✅ Tokenizer đã tải thành công!")
+        
+        # Kiểm tra xem có GPU không để quyết định precision
+        if torch.cuda.is_available():
+            print("Phát hiện GPU, sẽ sử dụng half precision...")
+            model = AutoModelForCausalLM.from_pretrained(
+                model_path,
+                torch_dtype=torch.float16,
+                device_map="auto"
+            )
+        else:
+            print("Không phát hiện GPU, sử dụng CPU...")
+            model = AutoModelForCausalLM.from_pretrained(model_path)
+        
+        print("✅ Model đã tải thành công!")
+        
+        # Kiểm tra pad token
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+            print("Đã thiết lập pad_token = eos_token")
+            
+        # Thử nghiệm mô hình với một câu đơn giản
+        test_text = "Xin chào, tôi là"
+        print(f"Thử nghiệm với input: '{test_text}'")
+        
+        inputs = tokenizer(test_text, return_tensors="pt")
+        if torch.cuda.is_available():
+            inputs = inputs.to("cuda")
+            
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=20,
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.95
+            )
+            
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        print(f"Kết quả: '{generated_text}'")
+        
+    except Exception as e:
+        print(f"Lỗi khi kiểm tra mô hình: {e}")
+    
+    return model_path
+
+if __name__ == "__main__":
+    download_model()