Update README.md

README.md

@@ -28,7 +28,7 @@ def chat_format(conversation:list,tokenizer,chat_type="mistral"):
         id = tokenizer.apply_chat_template(ap,chat_template=mistral_template,tokenize=False)
     elif chat_type=='llama3':
         id = tokenizer.apply_chat_template(ap,chat_template=llama3_template,tokenize=False)
-        id = id.rstrip("<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n")
+        #id = id.rstrip("<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n")
     return id
 
 user_chat=[{"role":"user","content":"In a basket, there are 20 oranges, 60 apples, and 40 bananas. If 15 pears were added, and half of the oranges were removed, what would be the new ratio of oranges to apples, bananas, and pears combined within the basket?"}]
@@ -48,12 +48,12 @@ def predict(content_prompt):
                     num_beams=1,
                     do_sample=True,
                     repetition_penalty=1.0,
-                    eos_token_id=tokenizer.eos_token_id,
-                    pad_token_id=tokenizer.pad_token_id,
+                    eos_token_id=tokenizer.convert_tokens_to_ids("<|eot_id|>"),
+                    pad_token_id=tokenizer.convert_tokens_to_ids("<|eot_id|>"),
                 )
         s = generation_output.sequences[0]
         output = tokenizer.decode(s,skip_special_tokens=True)
-        output1 = output.split("<|eot_id|>")[-1].strip()
+        output1 = output.split("<|eot_id|>")[-2].lstrip("<|start_header_id|>assistant<|end_header_id|>").strip()
         # print(output1)
     return output1