Update README.md

README.md

@@ -31,15 +31,17 @@ You can find more details in the [technical report](https://arxiv.org/abs/2310.1
 
 
 ## Use in Transformers
-## Load model directly
+**Load model directly** \
 import torch \
-from transformers import AutoTokenizer, AutoModelForCausalLM, MistralForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, MistralForCausalLM \
+from huggingface_hub import hf_hub_download
 
 model = MistralForCausalLM.from_pretrained("ai-agi/neural-zephyr", use_cache=False,  torch_dtype=torch.bfloat16, device_map="auto") \
-state_dict = torch.load('model_weights.pth') \
-model.load_state_dict(state_dict) 
+model_weights = hf_hub_download(repo_id="ai-agi/neural-zephyr", filename="model_weights.pth") \
+state_dict = torch.load(model_weights) \
+model.load_state_dict(state_dict)
 
 tokenizer = AutoTokenizer.from_pretrained("ai-agi/neural-zephyr", use_fast=True) \
 if tokenizer.pad_token is None: \
-    tokenizer.pad_token = tokenizer.eos_token)
-
+    tokenizer.pad_token = tokenizer.eos_token) \
+**Manage your GPU/CPU memory for model and weights**