Update app.py

app.py

@@ -1,49 +1,164 @@
+import os
+from threading import Thread
+from typing import Iterator, List, Dict, Any
+
 import gradio as gr
-from transformers import pipeline, Conversation, AutoTokenizer
+import spaces
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, Conversation, pipeline
+
+MAX_MAX_NEW_TOKENS = 1024
+DEFAULT_MAX_NEW_TOKENS = 256
+MAX_INPUT_TOKEN_LENGTH = 512
+
+DESCRIPTION = """\
+# Buzz-3B-Small
+This Space demonstrates Buzz-3b-small-v0.6.3.
+"""
+
+LICENSE = """
+<p/>
+---
+Chat with Buzz-small!
+only 3b,  this demo runs on the fp16 weights of the model in safetensors format (converting to cpp soon)
+"""
+
+if torch.cuda.is_available():
+    torch.device('cuda')
+else:
+    torch.device('cpu')
+
+model_id = "H-D-T/Buzz-3b-small-v0.6.3"
+chatbot = pipeline(model=model_id, device=device, task="conversational")
 
-# Define the special tokens
+tokenizer = AutoTokenizer.from_pretrained(model_id)
 bos_token = "<|begin_of_text|>"
 eos_token = "<|eot_id|>"
 start_header_id = "<|start_header_id|>"
 end_header_id = "<|end_header_id|>"
 
-# Load the conversational pipeline and tokenizer
-model_id = "H-D-T/Buzz-3b-small-v0.6.3"
-chatbot = pipeline("conversational", model=model_id)
-tokenizer = AutoTokenizer.from_pretrained(model_id)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.pad_token_id = tokenizer.eos_token_id
+    model.config.pad_token_id = tokenizer.eos_token_id
 
-def format_conversation(chat_history):
-    formatted_history = ""
-    for i, (user, assistant) in enumerate(chat_history):
-        user_msg = f"{start_header_id}user{end_header_id}\n\n{user.strip()}{eos_token}"
-        assistant_msg = f"{start_header_id}assistant{end_header_id}\n\n{assistant.strip()}{eos_token}"
+def format_conversation(chat_history: List[Dict[str, str]], add_generation_prompt=False) -> str:
+    """
+    Formats the chat history according to the model's chat template.
+    """
+    formatted_history = []
+    for i, message in enumerate(chat_history):
+        role, content = message["role"], message["content"]
+        formatted_message = f"{start_header_id}{role}{end_header_id}\n\n{content.strip()}{eos_token}"
         if i == 0:
-            user_msg = bos_token + user_msg
-        formatted_history += user_msg + assistant_msg
-    return formatted_history
-
-def predict(message, chat_history):
-    chat_history.append(("user", message))
-    formatted_history = format_conversation(chat_history)
-    conversation = Conversation(formatted_history)
-    conversation = chatbot(conversation)
-    response = conversation.generated_responses[-1]
-    chat_history.append(("assistant", response))
-    return "", chat_history
+            formatted_message = bos_token + formatted_message
+        formatted_history.append(formatted_message)
+    
+    if add_generation_prompt:
+        formatted_history.append(f"{start_header_id}assistant{end_header_id}\n\n")
+    else:
+        formatted_history.append(eos_token)
+    
+    return "".join(formatted_history)
+
+@spaces.GPU
+def generate(
+    message: str,
+    chat_history: List[Dict[str, str]],
+    max_new_tokens: int = 1024,
+    temperature: float = 0.6,
+    top_p: float = 0.9,
+    top_k: int = 50,
+    repetition_penalty: float = 1.4,
+) -> Iterator[str]:
+    
+    chat_history.append({"role": "user", "content": message})
+    chat_context = format_conversation(chat_history, add_generation_prompt=True)
+    input_ids = tokenizer([chat_context], return_tensors="pt").input_ids
+    
+    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
+        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
+        gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
+    input_ids = input_ids.to(device)
+    
+    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
+    generate_kwargs = dict(
+        input_ids=input_ids,
+        streamer=streamer,
+        max_new_tokens=max_new_tokens,
+        do_sample=True,
+        top_p=top_p,
+        top_k=top_k,
+        temperature=temperature,
+        num_beams=1,
+        pad_token_id=tokenizer.eos_token_id,
+        repetition_penalty=repetition_penalty,
+        no_repeat_ngram_size=5,
+        early_stopping=False,
+    )
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
+    
+    outputs = []
+    for text in streamer:
+        outputs.append(text)
+        yield "".join(outputs)
+
+
+chat_interface = gr.ChatInterface(
+    fn=generate,
+    additional_inputs=[
+        gr.Slider(
+            label="Max new tokens",
+            minimum=1,
+            maximum=MAX_MAX_NEW_TOKENS,
+            step=1,
+            value=DEFAULT_MAX_NEW_TOKENS,
+        ),
+        gr.Slider(
+            label="Temperature",
+            minimum=0.1,
+            maximum=4.0,
+            step=0.1,
+            value=0.6,
+        ),
+        gr.Slider(
+            label="Top-p (nucleus sampling)",
+            minimum=0.05,
+            maximum=1.0,
+            step=0.05,
+            value=0.9,
+        ),
+        gr.Slider(
+            label="Top-k",
+            minimum=1,
+            maximum=1000,
+            step=1,
+            value=50,
+        ),
+        gr.Slider(
+            label="Repetition penalty",
+            minimum=1.0,
+            maximum=2.0,
+            step=0.05,
+            value=1.4,
+        ),
+    ],
+    stop_btn=None,
+    examples=[
+        ["A recipe for a chocolate cake:"],
+        ["Can you explain briefly to me what is the Python programming language?"],
+        ["Explain the plot of Cinderella in a sentence."],
+        ["Question: What is the capital of France?\nAnswer:"],
+        ["Question: I am very tired, what should I do?\nAnswer:"],
+    ],
+)
 
 with gr.Blocks(css="style.css") as demo:
-    gr.Markdown("# Buzz-3B-Small Conversational Demo")
-    with gr.Chatbot() as chatbot_ui:
-        chatbot_ui.append({"role": "assistant", "content": "Hi, how can I help you today?"})
-    with gr.Row():
-        with gr.Column():
-            textbox = gr.Textbox(label="Your message:")
-        with gr.Column():
-            submit_btn = gr.Button("Send")
-
-    chat_history = gr.State([])
-
-    submit_btn.click(predict, inputs=[textbox, chat_history], outputs=[textbox, chat_history])
+    gr.Markdown(DESCRIPTION)
+    gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")
+    chat_interface.render()
+    gr.Markdown(LICENSE)
 
 if __name__ == "__main__":
     demo.queue(max_size=20).launch()