Update app.py

app.py

@@ -6,14 +6,9 @@ import gradio as gr
 import sentencepiece
 
 
-# Load the model and tokenizer
-model_name = "01-ai/Yi-34B-200K"
-model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
-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)
+# Load the model and tokenizer using transformers
+model = AutoModelForCausalLM.from_pretrained("01-ai/Yi-34B", device_map="auto", torch_dtype="auto", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("01-ai/Yi-34B", trust_remote_code=True)
 
 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)
@@ -58,7 +53,7 @@ MAX_MAX_NEW_TOKENS = 200000
 DEFAULT_MAX_NEW_TOKENS = 100000
 MAX_INPUT_TOKEN_LENGTH = 100000
 
-DESCRIPTION = "# [Yi-6B](https://huggingface.co/01-ai/Yi-6B)"
+DESCRIPTION = "# 👋🏻Welcome to 🙋🏻‍♂️Tonic's🧑🏻‍🚀YI-200K🚀"
 
 def clear_and_save_textbox(message): return '', message