Update README.md

README.md

@@ -65,12 +65,14 @@ import torch
 
 from Utils import GPTJBlock, GPTJForCausalLM
 
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 transformers.models.gptj.modeling_gptj.GPTJBlock = GPTJBlock  # monkey-patch GPT-J
 
-tokenizer = transformers.AutoTokenizer.from_pretrained("mrm8488/bertin-gpt-j-6B-ES-8bit")
-model = GPTJForCausalLM.from_pretrained("hivemind/gpt-j-6B-8bit", pad_token_id=tokenizer.eos_token_id, low_cpu_mem_usage=True).to(device)
+ckpt = "mrm8488/bertin-gpt-j-6B-ES-8bit"
+
+tokenizer = transformers.AutoTokenizer.from_pretrained(ckpt)
+model = GPTJForCausalLM.from_pretrained(ckpt, pad_token_id=tokenizer.eos_token_id, low_cpu_mem_usage=True).to(device)
 
 
 prompt = tokenizer("El sentido de la vida es", return_tensors='pt')