Update app.py

app.py

@@ -7,9 +7,9 @@ from tokenization_yi import YiTokenizer
 
 
 os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:120'
-model_id = "larryvrh/Yi-6B-200K-Llamafied"
+model_id = "01-ai/Yi-6B-200K"
 tokenizer_path = "./"
-eos_token_id = 7
+# eos_token_id = 7
 
 DESCRIPTION = """
 # 👋🏻Welcome to 🙋🏻‍♂️Tonic's🧑🏻‍🚀YI-200K🚀
@@ -18,18 +18,18 @@ You can also use 🧑🏻‍🚀YI-200K🚀 by cloning this space. 🧬🔬🔍
 Join us : 🌟TeamTonic🌟 is always making cool demos! Join our active builder's🛠️community on 👻Discord: [Discord](https://discord.gg/nXx5wbX9) On 🤗Huggingface: [TeamTonic](https://huggingface.co/TeamTonic) & [MultiTransformer](https://huggingface.co/MultiTransformer) On 🌐Github: [Polytonic](https://github.com/tonic-ai) & contribute to 🌟 [PolyGPT](https://github.com/tonic-ai/polygpt-alpha)
 """
 
-# tokenizer = AutoTokenizer.from_pretrained(model_id)
-tokenizer = YiTokenizer.from_pretrained(tokenizer_path)
+tokenizer = AutoTokenizer.from_pretrained(model_id, device_map="auto", trust_remote_code=True)
+# tokenizer = YiTokenizer.from_pretrained(tokenizer_path)
 model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=torch.bfloat16, trust_remote_code=True)
-tokenizer.eos_token_id = eos_token_id
-model.config.eos_token_id = eos_token_id
+# tokenizer.eos_token_id = eos_token_id
+# model.config.eos_token_id = eos_token_id
 
 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_dict = {"role": "user", "content": f"{system_message}{user_message}"}
+    prompt_dict = {"role": "user", "content": f"{system_message}\n\n{user_message}"}
     prompt = str(prompt_dict)
     return prompt
 
-def predict(message, system_message, max_new_tokens=4056, temperature=3.5, top_p=0.9, top_k=800, model_max_length = 32000, do_sample=False):
+def predict(message, system_message, max_new_tokens=4056, temperature=3.5, top_p=0.9, top_k=40, model_max_length = 32000, do_sample=False):
     formatted_prompt = format_prompt(message, system_message)
 
     input_ids = tokenizer.encode(formatted_prompt, return_tensors='pt')
@@ -41,7 +41,7 @@ def predict(message, system_message, max_new_tokens=4056, temperature=3.5, top_p
         temperature=temperature,  
         top_p=top_p,              
         top_k=top_k,              
-        no_repeat_ngram_size=5,
+        no_repeat_ngram_size=9,
         pad_token_id=tokenizer.eos_token_id,
         do_sample=do_sample
     )
@@ -65,7 +65,7 @@ with gr.Blocks(theme='ParityError/Anime') as demo:
         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)
+        top_k = gr.Slider(label='Top-K', minimum=1, maximum=1000, step=1, value=40)
         do_sample_checkbox = gr.Checkbox(label='Disable for faster inference', value=True)
 
     submit_button.click(