Update main.py

main.py

@@ -1,13 +1,20 @@
-from huggingface_hub import snapshot_download
-from pathlib import Path
-import os
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Убедитесь, что переменная окружения HF_TOKEN установлена
-token = os.getenv("HF_TOKEN")
-if not token:
-    raise ValueError("Hugging Face token not found in environment variables")
+# Load model and tokenizer
+model_name = "mistralai/Codestral-22B-v0.1"
+model = AutoModelForCausalLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-mistral_models_path = Path.home().joinpath('mistral_models', 'Mamba-Codestral-7B-v0.1')
-mistral_models_path.mkdir(parents=True, exist_ok=True)
+# No need to move model to GPU, default is CPU
+# model.to("cpu")  # This line can be omitted since it's already on CPU by default
 
-snapshot_download(repo_id="mistralai/Mamba-Codestral-7B-v0.1", allow_patterns=["params.json", "consolidated.safetensors", "tokenizer.model.v3"], local_dir=mistral_models_path)
+# Encode input tokens
+input_text = "Your input text here"
+tokens = tokenizer(input_text, return_tensors="pt").input_ids
+
+# Generate output
+generated_ids = model.generate(tokens, max_new_tokens=1000, do_sample=True)
+
+# Decode generated tokens
+result = tokenizer.decode(generated_ids[0].tolist(), skip_special_tokens=True)
+print(result)