Update README.md

README.md

@@ -26,13 +26,14 @@ import torch
 
 model_id = "UnicomLLM/Unichat-llama3-Chinese-8B"
 
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    torch_dtype=torch.bfloat16,
-    device_map="auto",
+pipeline = transformers.pipeline(
+    "text-generation",
+    model=model_id,
+    model_kwargs={"torch_dtype": torch.bfloat16},
+    device="cuda",
 )
 
+
 messages = [
     {"role": "system", "content": "You are a helpful assistant"},
     {"role": "user", "content": "Who are you?"},