Update README.md

README.md

@@ -48,11 +48,10 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 torch.random.manual_seed(0)
 
 model = AutoModelForCausalLM.from_pretrained(
-    "OEvortex/EMO-phi-128k",
-    device_map="cuda",
-    torch_dtype="auto",
-    trust_remote_code=True,
-    ignore_mismatched_sizes=True
+    "OEvortex/EMO-phi-128k", 
+    device_map="cuda", 
+    torch_dtype="auto", 
+    trust_remote_code=True, 
 )
 tokenizer = AutoTokenizer.from_pretrained("OEvortex/EMO-phi-128k")
 
@@ -61,16 +60,6 @@ messages = [
     {"role": "user", "content": "My best friend recently lost their parent to cancer after a long battle. They are understandably devastated and struggling with grief."},
 ]
 
-# Prepare the input for the pipeline
-formatted_messages = ""
-for message in messages:
-    formatted_messages += f"<|im_start|>{message['role']}\n{message['content']}<|im_end|>\n"
-
-# Optionally, add a generation prompt
-add_generation_prompt = True # Set this to True if you want to add a generation prompt
-if add_generation_prompt:
-    formatted_messages += "<|im_start|>assistant\n"
-
 pipe = pipeline(
     "text-generation",
     model=model,
@@ -78,13 +67,14 @@ pipe = pipeline(
 )
 
 generation_args = {
-    "max_new_tokens": 500,
+    "max_new_tokens": 2024,
     "return_full_text": False,
     "temperature": 0.6,
     "do_sample": True,
 }
 
-output = pipe(formatted_messages, **generation_args)
+output = pipe(messages, **generation_args)
 print(output[0]['generated_text'])
 
+
 ```