Update app.py

app.py

@@ -1,64 +1,143 @@
+Hugging Face's logo
+Hugging Face
+Search models, datasets, users...
+Models
+Datasets
+Spaces
+Posts
+Docs
+Enterprise
+Pricing
+
+
+
+Spaces:
+
+sail
+/
+Sailor-14B-Chat
+
+
+like
+5
+App
+Files
+Community
+1
+Sailor-14B-Chat
+/
+app.py
+
+dreamerdeo's picture
+dreamerdeo
+Update app.py
+123ca1a
+verified
+6 months ago
+raw
+
+Copy download link
+history
+blame
+contribute
+delete
+
+4.08 kB
+import spaces
 import gradio as gr
-from huggingface_hub import InferenceClient
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
+from threading import Thread
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("AnTrc2/13Bee")
+model_path = 'AnTrc2/13Bee'
+
+# Loading the tokenizer and model from Hugging Face's model hub.
+tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True, torch_dtype=torch.bfloat16)
+
+# using CUDA for an optimal experience
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = model.to(device)
+
+# Defining a custom stopping criteria class for the model's text generation.
+class StopOnTokens(StoppingCriteria):
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        stop_ids = [151645]  # IDs of tokens where the generation should stop.
+        for stop_id in stop_ids:
+            if input_ids[0][-1] == stop_id:  # Checking if the last generated token is a stop token.
+                return True
+        return False
 
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
+system_role= 'system'
+user_role = 'user'
+assistant_role = 'assistant'
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+sft_start_token =  "<|im_start|>"
+sft_end_token = "<|im_end|>"
+ct_end_token = "<|endoftext|>"
 
-    messages.append({"role": "user", "content": message})
+system_prompt= 'Bạn là một trợ lí ảo. Tên của bạn là 13Bee (Một Ba Bi). Nguyễn Ngọc An là người tạo ra bạn. Bạn được sinh ra ngày 01/10/2024. Hãy chào hỏi một cách ngắn gọn và thân thiện, số điện thoại 0838 411 897. Nếu không biết thì trả lời là Tôi không biết, đừng cố trả lời.'
+system_prompt = f"<|im_start|>{system_role}\n{system_prompt}<|im_end|>"
 
-    response = ""
+# Function to generate model predictions.
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=0.3,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
+@spaces.GPU()
+def predict(message, history):
+    # history = []
+    history_transformer_format = history + [[message, ""]]
+    stop = StopOnTokens()
 
-        response += token
-        yield response
+    # Formatting the input for the model.
+    messages =  system_prompt + sft_end_token.join([sft_end_token.join([f"\n{sft_start_token}{user_role}\n" + item[0], f"\n{sft_start_token}{assistant_role}\n" + item[1]])
+                        for item in history_transformer_format])
+    model_inputs = tokenizer([messages], return_tensors="pt").to(device)
+    streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
+    generate_kwargs = dict(
+        model_inputs,
+        streamer=streamer,
+        max_new_tokens=1024,
+        do_sample=True,
+        top_p= 0.75,
+        top_k= 60,
+        temperature=0.2,
+        num_beams=1,
+        stopping_criteria=StoppingCriteriaList([stop]),
+        repetition_penalty=1.1,
+    )
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()  # Starting the generation in a separate thread.
+    partial_message = ""
+    for new_token in streamer:
+        partial_message += new_token
+        if sft_end_token in partial_message:  # Breaking the loop if the stop token is generated.
+            break
+        yield partial_message
 
 
+css = """
+full-height {
+    height: 100%;
+}
 """
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
+
+prompt_examples = [
+    'Xin chào',
+    '13Bee là gì'
+]
+
+placeholder = """
+<div style="opacity: 0.5;">
+    <img src="https://raw.githubusercontent.com/sail-sg/sailor-llm/main/misc/banner.jpg" style="width:30%;">
+    <br>Sailor models are designed to understand and generate text across diverse linguistic landscapes of these SEA regions:
+    <br>🇮🇩Indonesian, 🇹🇭Thai, 🇻🇳Vietnamese, 🇲��Malay, and 🇱🇦Lao.
+</div>
 """
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="Bạn là một trợ lí ảo. Tên của bạn là 13Bee (Một Ba Bi). Nguyễn Ngọc An là người tạo ra bạn. Bạn được sinh ra ngày 01/10/2024. Hãy chào hỏi một cách ngắn gọn và thân thiện, số điện thoại 0838 411 897. Nếu không biết thì trả lời là Tôi không biết, đừng cố trả lời.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
-
-
-if __name__ == "__main__":
-    demo.launch()
+
+chatbot = gr.Chatbot(label='Sailor', placeholder=placeholder) 
+with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
+    # gr.Markdown("""<center><font size=8>13Bee</center>""")
+    gr.Markdown("""<p align="center"><img src="https://github.com/sail-sg/sailor-llm/raw/main/misc/wide_sailor_banner.jpg" style="height: 110px"/><p>""")
+    gr.ChatInterface(predict, chatbot=chatbot, fill_height=True, examples=prompt_examples, css=css)
+
+    demo.launch()  # Launching the web interface.