Update README.md

README.md

@@ -43,19 +43,41 @@ The model is intended for commercial and research use in English.
 
 Use the code below to get started with the model.
 
-```bibtex
+```python
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline
 
 model_name = "Deci/DeciLM-7B-instruct"
+
 device = "cuda" # for GPU usage or "cpu" for CPU usage
 
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, trust_remote_code=True).to(device)
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit = True,
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
 
-inputs = tokenizer.encode("In a shocking finding, scientists discovered a herd of unicorns living in", return_tensors="pt").to(device)
-outputs = model.generate(inputs, max_new_tokens=100, do_sample=True, top_p=0.95)
-print(tokenizer.decode(outputs[0]))
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    device_map="auto",
+    trust_remote_code=True,
+    quantization_config=bnb_config
+)
+
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+tokenizer.pad_token = tokenizer.eos_token
+
+deci_generator = pipeline("text-generation",
+                          model=model,
+                          tokenizer=tokenizer,
+                          temperature=0.1,
+                          device_map="auto",
+                          max_length=4096,
+                          return_full_text=False
+)
+
+prompt = "How do I make the most delicious pancakes the world has ever tasted?"
+response = deci_generator(prompt)[0]['generated_text']
+print(response)
 ```
 
 ## Evaluation