Update app.py

app.py

@@ -9,6 +9,7 @@ from tokenization_yi import YiTokenizer
 os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:120'
 model_id = "larryvrh/Yi-6B-200K-Llamafied"
 tokenizer_path = "./"
+eos_token_id = 7
 
 DESCRIPTION = """
 # 👋🏻Welcome to 🙋🏻‍♂️Tonic's🧑🏻‍🚀YI-200K🚀"
@@ -19,13 +20,18 @@ Join us : 🌟TeamTonic🌟 is always making cool demos! Join our active builder
 
 tokenizer = AutoModelForCausalLM.from_pretrained(tokenizer_path)
 tokenizer = YiTokenizer.from_pretrained(tokenizer_path)
-model = AutoModelForCausalLM.from_pretrained(model_id=model_id, device_map="auto", torch_dtype="auto", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(model_id=tokenizer_path, device_map="auto", torch_dtype="auto", trust_remote_code=True)
 
-def predict(message, max_new_tokens=4056, temperature=3.5, top_p=0.9, top_k=800, do_sample=False):
+def format_prompt(user_message, system_message="I am YiTonic, an AI language model created by Tonic-AI. I am a cautious assistant. I carefully follow instructions. I am helpful and harmless and I follow ethical guidelines and promote positive behavior."):
+    prompt = f"<|im_start|>assistant\n{self.system_message}<|im_end|>\n<|im_start|>\nuser\n{user_message}<|im_end|>\nassistant\n"
+    return prompt
 
-    prompt = message.strip()
-    input_ids = tokenizer.encode(prompt, return_tensors='pt')
+def predict(message, system_message, max_new_tokens=4056, temperature=3.5, top_p=0.9, top_k=800, do_sample=False):
+    formatted_prompt = format_prompt(message, system_message)
+
+    input_ids = tokenizer.encode(formatted_prompt, return_tensors='pt')
     input_ids = input_ids.to(model.device)
+
     response_ids = model.generate(
         input_ids,
         max_length=max_new_tokens + input_ids.shape[1],
@@ -35,28 +41,33 @@ def predict(message, max_new_tokens=4056, temperature=3.5, top_p=0.9, top_k=800,
         pad_token_id=tokenizer.eos_token_id,
         do_sample=do_sample
     )
+
     response = tokenizer.decode(response_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
-    return [("bot", response)]
-    
+    return [("bot", response)]    
 
 with gr.Blocks(theme='ParityError/Anime') as demo:
     gr.Markdown(DESCRIPTION)
     with gr.Group():
-        textbox = gr.Textbox(placeholder='Enter your message here', label='Your Message', lines=2)
+        textbox = gr.Textbox(placeholder='Your Message Here', label='Your Message', lines=2)
+        system_prompt = gr.Textbox(placeholder='Provide a System Prompt In The First Person', label='System Prompt', lines=2, value="You are YiTonic, an AI language model created by Tonic-AI. You are a cautious assistant. You carefully follow instructions. You are helpful and harmless and you follow ethical guidelines and promote positive behavior.")
+
+    with gr.Group():
         submit_button = gr.Button('Submit', variant='primary')
-        chatbot = gr.Chatbot(label='TonicYi-6B-200K')
 
     with gr.Accordion(label='Advanced options', open=False):
-        max_new_tokens = gr.Slider(label='Max New Tokens', minimum=1, maximum=55000, step=1, value=8000)
+        max_new_tokens = gr.Slider(label='Max New Tokens', minimum=1, maximum=55000, step=1, value=4056)
         temperature = gr.Slider(label='Temperature', minimum=0.1, maximum=4.0, step=0.1, value=1.2)
         top_p = gr.Slider(label='Top-P (nucleus sampling)', minimum=0.05, maximum=1.0, step=0.05, value=0.9)
         top_k = gr.Slider(label='Top-K', minimum=1, maximum=1000, step=1, value=900)
-        do_sample_checkbox = gr.Checkbox(label='Disable for faster inference', value=False )
+        do_sample_checkbox = gr.Checkbox(label='Disable for faster inference', value=False)
 
     submit_button.click(
         fn=predict,
-        inputs=[textbox, max_new_tokens, temperature, top_p, top_k, do_sample_checkbox],
+        inputs=[textbox, system_prompt, max_new_tokens, temperature, top_p, top_k, do_sample_checkbox],
         outputs=chatbot
     )
 
+    with gr.Group():
+        chatbot = gr.Chatbot(label='TonicYi-6B-200K')
+
 demo.launch()