Update app.py

app.py

@@ -1,4 +1,5 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from optimum.bettertransformer import BetterTransformer
 from tokenization_yi import YiTokenizer
 import torch
 import os
@@ -6,13 +7,26 @@ import gradio as gr
 import sentencepiece
 
 
+DESCRIPTION = """
+# 👋🏻Welcome to 🙋🏻‍♂️Tonic's🧑🏻‍🚀YI-6B-200K🚀"
+You can use this Space to test out the current model [01-ai/Yi-6B-200K](https://huggingface.co/01-ai/Yi-6B-200K)
+You can also use 🧑🏻‍🚀YI-200K🚀 by cloning this space. 🧬🔬🔍 Simply click here: <a style="display:inline-block" href="https://huggingface.co/spaces/Tonic1/YiTonic?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a></h3> 
+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)
+"""
+
+MAX_MAX_NEW_TOKENS = 160000
+DEFAULT_MAX_NEW_TOKENS = 20000
+MAX_INPUT_TOKEN_LENGTH = 160000
+
+
 os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:126'
 
 # Load the model and tokenizer using transformers
 model = AutoModelForCausalLM.from_pretrained("01-ai/Yi-6B-200K", device_map="auto", torch_dtype="auto", trust_remote_code=True)
 tokenizer = YiTokenizer(vocab_file="./tokenizer.model")
+model = BetterTransformer.transform(model)
 
-def run(message, chat_history, max_new_tokens=4056, temperature=3.5, top_p=0.9, top_k=800):
+def run(message, chat_history, max_new_tokens=20000, temperature=1.0, top_p=0.9, top_k=0.9):
     prompt = get_prompt(message, chat_history)
 
     # Encode the prompt to tensor
@@ -49,17 +63,6 @@ def get_prompt(message, chat_history):
     texts.append(f"{message}")
     return ''.join(texts)
 
-DESCRIPTION = """
-# 👋🏻Welcome to 🙋🏻‍♂️Tonic's🧑🏻‍🚀YI-6B-200K🚀"
-You can use this Space to test out the current model [01-ai/Yi-6B-200K](https://huggingface.co/01-ai/Yi-6B-200K)
-You can also use 🧑🏻‍🚀YI-200K🚀 by cloning this space. 🧬🔬🔍 Simply click here: <a style="display:inline-block" href="https://huggingface.co/spaces/Tonic1/YiTonic?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a></h3> 
-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)
-"""
-
-MAX_MAX_NEW_TOKENS = 16000
-DEFAULT_MAX_NEW_TOKENS = 4056
-MAX_INPUT_TOKEN_LENGTH = 32000
-
 def clear_and_save_textbox(message): return '', message
 
 def display_input(message, history=[]):
@@ -83,7 +86,7 @@ def generate(message, history_with_input, max_new_tokens, temperature, top_p, to
 
 
 def process_example(message):
-    generator = generate(message, [], 1024, 2.5, 0.95, 900)
+    generator = generate(message, [], 4056, 1.9, 0.95, 900)
     for x in generator:
         pass
     return '', x
@@ -119,7 +122,7 @@ with gr.Blocks(theme='ParityError/Anime') as demo:
     with gr.Accordion(label='Advanced options', open=False):
 #       system_prompt = gr.Textbox(label='System prompt', value=DEFAULT_SYSTEM_PROMPT, lines=5, interactive=False)
         max_new_tokens = gr.Slider(label='Max New Tokens', minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1, value=DEFAULT_MAX_NEW_TOKENS)
-        temperature = gr.Slider(label='Temperature', minimum=0.1, maximum=4.0, step=0.1, value=0.1)
+        temperature = gr.Slider(label='Temperature', minimum=0.1, maximum=2.0, step=0.1, value=0.1)
         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=10)