Update README.md

README.md

@@ -11,10 +11,21 @@ pip install git+https://huggingface.co/CognitivessAI/cognitivess
 Then, you can use the model like this:
 
 ```python
-# pip install bitsandbytes accelerate
+# Install required packages
+#pip install bitsandbytes accelerate
+#pip install git+https://huggingface.co/CognitivessAI/cognitivess
+
+# Import necessary libraries
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import torch
 
+# Import and register your custom classes
+from cognitivess_model import CognitivessConfig, CognitivessForCausalLM
+from transformers import AutoConfig, AutoModelForCausalLM
+
+AutoConfig.register("cognitivess", CognitivessConfig)
+AutoModelForCausalLM.register(CognitivessConfig, CognitivessForCausalLM)
+
 # Set up quantization config
 quantization_config = BitsAndBytesConfig(load_in_8bit=True)
 
@@ -23,12 +34,12 @@ 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
+    device_map="auto"
 )
 
 # Prepare input
 input_text = "Write me a poem about Machine Learning."
-inputs = tokenizer(input_text, return_tensors="pt").to("cuda")
+inputs = tokenizer(input_text, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
 
 # Generate output
 with torch.no_grad():