Update README.md

README.md

@@ -118,38 +118,50 @@ To use this model, first install the custom package:
 Then, you can use the model like this:
 
 ```python
-import cognitivess_model  # Ensure this imports the custom model package
-
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import cognitivess_model
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-model_id = "CognitivessAI/cognitivess"
+# Define the model path
+model_path = "CognitivessAI/cognitivess"
 
 # Load the tokenizer
-tokenizer = AutoTokenizer.from_pretrained(model_id)
+tokenizer = AutoTokenizer.from_pretrained(model_path)
 
-# Load the model
+# Load the model with correct configuration for precision and device placement
 model = AutoModelForCausalLM.from_pretrained(
-    model_id,
+    model_path,
     torch_dtype=torch.float32,
-    device_map="auto"
-)
+    device_map="auto"  # Automatically maps model to available devices
+).eval()
 
+# Move model to CUDA if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
+# Prepare input
 messages = [
-    {"role": "user", "content": "Explain how large language models work in detail."},
+    {"role": "user", "content": "Who are you?"}
 ]
 
-input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
+# Tokenize input
+input_ids = tokenizer(
+    [msg["content"] for msg in messages], 
+    return_tensors='pt', 
+    padding=True, 
+    truncation=True
+).input_ids
 
-outputs = model.generate(
-    input_ids,
-    do_sample=True,
-    temperature=0.5,
-    max_new_tokens=1024
-)
+# Move input_ids to the same device as the model
+input_ids = input_ids.to(device)
 
-response = outputs[0][input_ids.shape[-1]:]
-print(tokenizer.decode(response, skip_special_tokens=True))
+# Generate output
+output_ids = model.generate(input_ids, max_new_tokens=50)
+
+# Decode output
+response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+print(response)
 
 ```