Update app.py

app.py

@@ -3,6 +3,8 @@ import torch
 import os
 import gradio as gr
 import sentencepiece
+from tokenization_yi import YiTokenizer
+
 
 model_id = "TheBloke/Yi-34B-200K-Llamafied-GPTQ"
 
@@ -11,7 +13,7 @@ gptq_config = GPTQConfig(
     exllama_config={"version": 2},
     disable_exllama=True
 )
-tokenizer = AutoTokenizer.from_pretrained(model_id)
+tokenizer = YiTokenizer.from_pretrained("./")
 model = AutoModelForCausalLM.from_pretrained(
     model_id, 
     device_map="auto",