Update README.md

README.md

@@ -41,13 +41,12 @@ Use the code below to get started with the model.
 from transformers import  AutoModelForCausalLM, AutoTokenizer 
 import torch
 
-model =  AutoModelForCausalLM.from_pretrained("DIAG-PSSeng/cicero_v2-phi1.5", torch_dtype=torch.float16).to("cuda")
+model =  AutoModelForCausalLM.from_pretrained("DIAG-PSSeng/cicero_v2-phi1.5", trust_remote_code=True ,torch_dtype=torch.float16).to("cuda")
 
 tokenizer = AutoTokenizer.from_pretrained("DIAG-PSSeng/cicero_v2-phi1.5", trust_remote_code=True)
 
 def generate_text(model, tokenizer, prompt, length=50, do_sample=True):
-    device="cuda"
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
     gen_tokens = model.generate(**inputs,do_sample=True,temperature=0.9, min_length=length,max_length=length)
     generated_text = tokenizer.batch_decode(gen_tokens)
     return generated_text