Update README.md

README.md

@@ -11,8 +11,29 @@ pip install git+https://huggingface.co/CognitivessAI/cognitivess
 Then, you can use the model like this:
 
 ```python
-from transformers import AutoTokenizer, AutoModelForCausalLM
+# pip install bitsandbytes accelerate
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+import torch
 
-tokenizer = AutoTokenizer.from_pretrained('CognitivessAI/cognitivess')
-model = AutoModelForCausalLM.from_pretrained('CognitivessAI/cognitivess')
+# Set up quantization config
+quantization_config = BitsAndBytesConfig(load_in_8bit=True)
+
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained("CognitivessAI/cognitivess")
+model = AutoModelForCausalLM.from_pretrained(
+    "CognitivessAI/cognitivess",
+    quantization_config=quantization_config,
+    device_map="auto"  # This will automatically distribute the model across available GPUs
+)
+
+# Prepare input
+input_text = "Write me a poem about Machine Learning."
+inputs = tokenizer(input_text, return_tensors="pt").to("cuda")
+
+# Generate output
+with torch.no_grad():
+    outputs = model.generate(**inputs, max_length=100)
+
+# Decode and print the result
+print(tokenizer.decode(outputs[0], skip_special_tokens=True))
 ```