adapted the safetensor example since injecting the system prompt manually is no longer necessary

README.md

@@ -23,31 +23,24 @@ A finetune of https://huggingface.co/stabilityai/stable-code-instruct-3b trained
 ### Safetensors (recommended)
 
 ```python
-from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Load the model and tokenizer
-print("Loading model...")
-model_name = "path/to/your/safetensors/model" #./stable-cypher-instruct-3b
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
+tokenizer = AutoTokenizer.from_pretrained("lakkeo/stable-cypher-instruct-3b", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained("lakkeo/stable-cypher-instruct-3b", torch_dtype=torch.bfloat16, trust_remote_code=True)
 
-# Define your question
-'''instruction is VERY IMPORTANT the model was finetuned on this particular system prompt. 
-Except bad performance without it'''
-instruction = "Create a Cypher statement to answer the following question:"
-question = "List the first 3 articles mentioning organizations with a revenue less than 5 million."
+messages = [
+    {
+        "role": "user",
+        "content": "Show me the people who have Python and Cloud skills and have been in the company for at least 3 years."
+    }
+]
 
-# Create the full prompt
-full_prompt = f"{instruction}\n\nHuman: {question}\n\nAssistant:"
+prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
 
-# Tokenize input
-inputs = tokenizer(full_prompt, return_tensors="pt")
+inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
 
-# Generate response
-print("Generating response...")
-with torch.no_grad():
-    outputs = model.generate(
+tokens = model.generate(
         **inputs,
         max_new_tokens=128,
         do_sample=True,
@@ -56,13 +49,7 @@ with torch.no_grad():
         pad_token_id=tokenizer.eos_token_id,
     )
 
-# Decode and print the generated response
-answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
-answer = answer[len(full_prompt):].strip()  # Remove the input prompt from the output
-
-print("\nQuestion:", question)
-print("\nGenerated Cypher statement:")
-print(answer)
+outputs = tokenizer.batch_decode(tokens[:, inputs.input_ids.shape[-1]:], skip_special_tokens=False)[0]
 ```
 
 ### GGUF