Update app.py

app.py

@@ -1,10 +1,15 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from tokenization_yi import YiTokenizer
 import torch
 
 # Load the model and tokenizer
 model_name = "01-ai/Yi-34B-200K"
 model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+tokenizer_dir = "Tonic1/YiTonic"
+vocab_file = os.path.join(tokenizer_dir, "tokenizer.model")
+tokenizer_json = os.path.join(tokenizer_dir, "tokenizer.json")
+tokenizer_config = os.path.join(tokenizer_dir, "tokenizer_config.json")
+tokenizer = YiTokenizer(vocab_file=vocab_file)
 
 def run(message, chat_history, system_prompt, max_new_tokens=1024, temperature=0.3, top_p=0.9, top_k=50):
     prompt = get_prompt(message, chat_history, system_prompt)