Update README.md

README.md

@@ -37,7 +37,7 @@ More information needed
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
 import torch
 
-MODEL_NAME = "ce-lery/mistral-2b-base"
+model_path = "ce-lery/mistral-2b-base"
 torch.set_float32_matmul_precision('high')
 
 device = "cuda"
@@ -49,12 +49,11 @@ model = AutoModelForCausalLM.from_pretrained(model_path,
                                              trust_remote_code=True,
                                              ).to(device)
 
-
 prompt = "自然言語処理とは、"
-
-inputs = tokenizer(prompt, add_special_tokens=True,return_tensors="pt").to(model.device)
+inputs = tokenizer(prompt,
+                   add_special_tokens=True,
+                   return_tensors="pt").to(model.device)
 with torch.no_grad():
-
     outputs = model.generate(
         inputs["input_ids"],
         max_new_tokens=4096,