Update app.py

app.py

@@ -1,4 +1,4 @@
-from transformers import AutoModelForCausalLM, GPTQConfig, AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import os
 import gradio as gr
@@ -6,15 +6,8 @@ import sentencepiece
 # from tokenization_yi import YiTokenizer
 
 
-from transformers import AutoModelForCausalLM, GPTQConfig, AutoTokenizer, AutoModelForCausalLM
-import torch
-import os
-import gradio as gr
-import sentencepiece
-from tokenization_yi import YiTokenizer
-
-os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:50'
-model_id = "TheBloke/Yi-34B-200K-Llamafied-GPTQ"
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:120'
+model_id = "larryvrh/Yi-34B-200K-Llamafied"
 
 
 tokenizer = AutoTokenizer.from_pretrained("larryvrh/Yi-34B-200K-Llamafied")
@@ -54,7 +47,7 @@ with gr.Blocks(theme='ParityError/Anime') as demo:
         chatbot = gr.Chatbot(label='TonicYi-30B-200K')
 
     with gr.Accordion(label='Advanced options', open=False):
-        max_new_tokens = gr.Slider(label='Max New Tokens', minimum=1, maximum=2056, step=1, value=980)
+        max_new_tokens = gr.Slider(label='Max New Tokens', minimum=1, maximum=55000, step=1, value=3800)
         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)