Upload README.md

README.md

@@ -53,24 +53,17 @@ dtype: bfloat16
 ## 💻 Usage
 
 ```python
-!pip install -qU transformers accelerate
-
-from transformers import AutoTokenizer
-import transformers
 import torch
-
-model = "vince62s/Phi-2-psy"
-messages = [{"role": "user", "content": "What is a large language model?"}]
-
-tokenizer = AutoTokenizer.from_pretrained(model)
-prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-pipeline = transformers.pipeline(
-    "text-generation",
-    model=model,
-    torch_dtype=torch.float16,
-    device_map="auto",
-)
-
-outputs = pipeline(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
-print(outputs[0]["generated_text"])
+from transformers import AutoModelForCausalLM, AutoTokenizer
+torch.set_default_device("cuda")
+model = AutoModelForCausalLM.from_pretrained("vince62s/phi-2-psy", torch_dtype="auto", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("vince62s/phi-2-psy", trust_remote_code=True)
+inputs = tokenizer('''def print_prime(n):
+   """
+   Print all primes between 1 and n
+   """''', return_tensors="pt", return_attention_mask=False)
+outputs = model.generate(**inputs, max_length=200)
+text = tokenizer.batch_decode(outputs)[0]
+print(text)
 ```
+