Update main.py

main.py

@@ -1,12 +1,13 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from huggingface_hub import snapshot_download
+from pathlib import Path
+import os
 
-model_id = "mistralai/Codestral-22B-v0.1"
-tokenizer = AutoTokenizer.from_pretrained(model_id)
+# Убедитесь, что переменная окружения HF_TOKEN установлена
+token = os.getenv("HF_TOKEN")
+if not token:
+    raise ValueError("Hugging Face token not found in environment variables")
 
-model = AutoModelForCausalLM.from_pretrained(model_id)
+mistral_models_path = Path.home().joinpath('mistral_models', 'Mamba-Codestral-7B-v0.1')
+mistral_models_path.mkdir(parents=True, exist_ok=True)
 
-text = "Hello my name is"
-inputs = tokenizer(text, return_tensors="pt")
-
-outputs = model.generate(**inputs, max_new_tokens=20)
-print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+snapshot_download(repo_id="mistralai/Mamba-Codestral-7B-v0.1", allow_patterns=["params.json", "consolidated.safetensors", "tokenizer.model.v3"], local_dir=mistral_models_path)