Update app.py

app.py

@@ -1,117 +1,79 @@
-import os
-import math
-import transformers
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+import os
 import gradio as gr
 import sentencepiece
-import gc
-
-title = "# Welcome to 🙋🏻‍♂️Tonic's🌷Tulu Chat!"
-description = """[allenai/tulu-2-dpo-7b](https://huggingface.co/allenai/tulu-2-dpo-7b) and larger Tulu-2 models are Instruct Llama Finetunes using the [mistralai/Mistral-7B](https://huggingface.co/mistralai/Mistral-7B-v0.1) recipe. You can use [allenai/tulu-2-13b](https://huggingface.co/allenai/tulu-2-13b) here via API using Gradio by scrolling down and clicking Use 'Via API' or privately by [cloning this space on huggingface](https://huggingface.co/spaces/Tonic1/TuluDemo?duplicate=true) See also the large model here : [allenai/tulu-2-dpo-70b](https://huggingface.co/allenai/tulu-2-dpo-70b) . [Join my active builders' server on discord](https://discord.gg/VqTxc76K3u). Let's build together!. [Add this Space as a discord bot to your server by clicking this link](https://discord.com/oauth2/authorize?client_id=1176628808212828231&scope=bot+applications.commands&permissions=326417525824). Big thanks to 🤗Huggingface Organisation for the🫂Community Grant"""
-
-os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:50'
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-model_name = "allenai/tulu-2-dpo-13b"
-tokenizer = AutoTokenizer.from_pretrained("allenai/tulu-2-dpo-13b")
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto")
-
-# bos_token_id = 1  
-# eos_token_id = 2  
-# tokenizer.bos_token_id = bos_token_id
-# tokenizer.eos_token_id = eos_token_id
-# model.config.bos_token_id = bos_token_id
-# model.config.eos_token_id = eos_token_id
-# if tokenizer.pad_token is None:
-#     tokenizer.pad_token = tokenizer.eos_token
-# model.config.pad_token_id = tokenizer.convert_tokens_to_ids(tokenizer.pad_token)
-
-class TuluChatBot:
-    def __init__(self, model, tokenizer, system_message="You are 🌷Tulu, an AI language model created by Tonic-AI. You are a cautious assistant. You carefully follow instructions. You are helpful and harmless and you follow ethical guidelines and promote positive behavior."):
-        self.model = model
-        self.tokenizer = tokenizer
-        self.system_message = system_message
-
-    def set_system_message(self, new_system_message):
-        self.system_message = new_system_message
-
-    def format_prompt(self, user_message):
-        prompt = f"<|assistant|>\n{self.system_message}\n<|user|>{user_message}\n<|assistant|>\n"
-        return prompt
-
-    def Tulu(self, user_message, temperature, max_new_tokens, top_p, repetition_penalty, do_sample):
-        try:
-            prompt = self.format_prompt(user_message)
-            inputs = self.tokenizer(prompt, return_tensors='pt', add_special_tokens=True)
-            input_ids = inputs["input_ids"].to(self.model.device)
-            attention_mask = inputs["attention_mask"].to(self.model.device)
-
-
-            output_ids = self.model.generate(
-                input_ids,
-                attention_mask=attention_mask,
-                max_length=input_ids.shape[1] + max_new_tokens,
-                temperature=temperature,
-                top_p=top_p,
-                repetition_penalty=repetition_penalty,
-                do_sample=do_sample
-            )
-
-            response = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
-            response = response.strip()
-            response = response.split("<|assistant|>\n")[-1]
-            return response        
-        finally:
-            del input_ids, attention_mask, output_ids
-            gc.collect()
-            torch.cuda.empty_cache()
-
-def gradio_Tulu(user_message, system_message, max_new_tokens, temperature, top_p, repetition_penalty, do_sample):
-    Tulu_bot.set_system_message(system_message)
-    if not do_sample:
-        max_length = 780
-        temperature = 0.9  
-        top_p = 0.9        
-        repetition_penalty = 0.9  
-    response = Tulu_bot.Tulu(user_message, temperature, max_new_tokens, top_p, repetition_penalty, do_sample)
-    return response
-
-    # Initialize TuluChatBot
-Tulu_bot = TuluChatBot(model, tokenizer)
+from tokenization_yi import YiTokenizer
+
+
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:120'
+model_id = "01-ai/Yi-6B-200K"
+tokenizer_path = "./"
+eos_token_id = 7
+
+DESCRIPTION = """
+# 👋🏻Welcome to 🙋🏻‍♂️Tonic's🧑🏻‍🚀YI-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) "🦙Llamified" version based on [01-ai/Yi-34B](https://huggingface.co/01-ai/Yi-34B)
+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)
+"""
+
+tokenizer = AutoTokenizer.from_pretrained(model_id, device_map="auto", trust_remote_code=True)
+# tokenizer = YiTokenizer.from_pretrained(tokenizer_path)
+model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=torch.bfloat16, trust_remote_code=True)
+tokenizer.eos_token_id = eos_token_id
+model.config.eos_token_id = eos_token_id
+
+def format_prompt(user_message, system_message="You are YiTonic, an AI language model created by Tonic-AI. You are a cautious assistant. You carefully follow instructions. You are helpful and harmless and follow ethical guidelines and promote positive behavior."):
+    prompt = f"<|im_start|>assistant\n{system_message}<|im_end|>\n<|im_start|>\nuser\n{user_message}<|im_end|>\nassistant\n"
+    return prompt
+
+def predict(message, system_message, max_new_tokens=4056, temperature=3.5, top_p=0.9, top_k=40, model_max_length = 32000, do_sample=False):
+    formatted_prompt = format_prompt(message, system_message)
+
+    input_ids = tokenizer.encode(formatted_prompt, return_tensors='pt')
+    input_ids = input_ids.to(model.device)
+
+    response_ids = model.generate(
+        input_ids,
+        max_length=max_new_tokens + input_ids.shape[1],
+        temperature=temperature,  
+        top_p=top_p,              
+        top_k=top_k,              
+        no_repeat_ngram_size=9,
+        pad_token_id=tokenizer.eos_token_id,
+        do_sample=do_sample
+    )
 
-# Gradio interface function
-def gradio_Tulu(user_message, system_message, max_new_tokens, temperature, top_p, repetition_penalty, do_sample):
-    Tulu_bot.set_system_message(system_message)
-    response = Tulu_bot.Tulu(user_message, temperature, max_new_tokens, top_p, repetition_penalty, do_sample)
-    return response
+    response = tokenizer.decode(response_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
+    truncate_str = "<|im_end|>"
+    if truncate_str and truncate_str in response:
+        response = response.split(truncate_str)[0]
 
+    return [("bot", response)]    
+with gr.Blocks(theme='ParityError/Anime') as demo:
+    gr.Markdown(DESCRIPTION)
+    with gr.Group():
+        textbox = gr.Textbox(placeholder='Your Message Here', label='Your Message', lines=2)
+        system_prompt = gr.Textbox(placeholder='Provide a System Prompt In The First Person', label='System Prompt', lines=2, value="You are YiTonic, an AI language model created by Tonic-AI. You are a cautious assistant. You carefully follow instructions. You are helpful and harmless and you follow ethical guidelines and promote positive behavior.")
 
-with gr.Blocks(theme = "ParityError/Anime") as demo:
-    gr.Markdown(title)  
-    gr.Markdown(description)  
-    with gr.Row():
-        system_message = gr.Textbox(label="Optional 🌷Tulu Assistant Message", lines=2)
-        user_message = gr.Textbox(label="Your Message", lines=3)
-    with gr.Row():
-        do_sample = gr.Checkbox(label="Advanced", value=True)
-    
-    with gr.Accordion("Advanced Settings", open=lambda do_sample: do_sample):
-        with gr.Row():
-            max_new_tokens = gr.Slider(label="Max new tokens", value=250, minimum=20, maximum=450, step=1)
-            temperature = gr.Slider(label="Temperature", value=0.3, minimum=0.1, maximum=1.0, step=0.1)
-            top_p = gr.Slider(label="Top-p (nucleus sampling)", value=0.90, minimum=0.01, maximum=0.99, step=0.05)
-            repetition_penalty = gr.Slider(label="Repetition penalty", value=0.9, minimum=0.05, maximum=1.0, step=0.05)
+    with gr.Group():
+        chatbot = gr.Chatbot(label='TonicYi-6B-200K-🧠🤯')
 
-    submit_button = gr.Button("Submit")
-    output_text = gr.Textbox(label="🌷Tulu Response")
+    with gr.Group():
+        submit_button = gr.Button('Submit', variant='primary')
 
-    def process(user_message, system_message, max_new_tokens, temperature, top_p, repetition_penalty, do_sample):
-        return gradio_Tulu(user_message, system_message, max_new_tokens, temperature, top_p, repetition_penalty, do_sample)
+    with gr.Accordion(label='Advanced options', open=False):
+        max_new_tokens = gr.Slider(label='Max New Tokens', minimum=1, maximum=55000, step=1, value=4056)
+        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=40)
+        do_sample_checkbox = gr.Checkbox(label='Disable for faster inference', value=True)
 
     submit_button.click(
-        process,
-        inputs=[user_message, system_message, max_new_tokens, temperature, top_p, repetition_penalty, do_sample],
-        outputs=output_text
+        fn=predict,
+        inputs=[textbox, system_prompt, max_new_tokens, temperature, top_p, top_k, do_sample_checkbox],
+        outputs=chatbot
     )
-    
+
 demo.launch()