Update app.py

app.py

@@ -5,48 +5,30 @@ import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
-MAX_INPUT_LIMIT = 3584
-MAX_NEW_TOKENS = 1536
-MODEL_NAME = "Azure99/blossom-v5.1-9b"
+MAX_NEW_TOKENS = 2048
+MODEL_NAME = "Azure99/Blossom-V6-14B"
 
 model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.bfloat16, device_map="auto")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
 
 def get_input_ids(inst, history):
-    prefix = ("A chat between a human and an artificial intelligence bot. "
-              "The bot gives helpful, detailed, and polite answers to the human's questions.")
-    patterns = []
-    for conv in history:
-        patterns.append(f'\n|Human|: {conv[0]}\n|Bot|: ')
-        patterns.append(f'{conv[1]}')
-    patterns.append(f'\n|Human|: {inst}\n|Bot|: ')
-    patterns[0] = prefix + patterns[0]
-
-    input_ids = []
-    for i, pattern in enumerate(patterns):
-        input_ids += tokenizer.encode(pattern, add_special_tokens=(i == 0))
-        if i % 2 == 1:
-            input_ids += [tokenizer.eos_token_id]
-    return input_ids
-
-
-def generate(generation_kwargs):
-    with torch.no_grad():
-        Thread(target=model.generate, kwargs=generation_kwargs).start()
+    conversation = []
+    for user, assistant in history:
+        conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
+    conversation.append({"role": "user", "content": inst})
+    return tokenizer.apply_chat_template(conversation, return_tensors="pt").to(model.device)
 
 
 @spaces.GPU
 def chat(inst, history, temperature, top_p, repetition_penalty):
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     input_ids = get_input_ids(inst, history)
-    if len(input_ids) > MAX_INPUT_LIMIT:
-        yield "The input is too long, please clear the history."
-        return
-    generation_kwargs = dict(input_ids=torch.tensor([input_ids]).to(model.device),
+    generation_kwargs = dict(input_ids=input_ids,
                              streamer=streamer, do_sample=True, max_new_tokens=MAX_NEW_TOKENS,
                              temperature=temperature, top_p=top_p, repetition_penalty=repetition_penalty)
-    generate(generation_kwargs)
+
+    Thread(target=model.generate, kwargs=generation_kwargs).start()
 
     outputs = ""
     for new_text in streamer:
@@ -87,7 +69,7 @@ additional_inputs = [
 gr.ChatInterface(chat,
                  chatbot=gr.Chatbot(show_label=False, height=500, show_copy_button=True, render_markdown=True),
                  textbox=gr.Textbox(placeholder="", container=False, scale=7),
-                 title="Blossom 9B Demo",
+                 title="Blossom 14B Demo",
                  description='Hello, I am Blossom, an open source conversational large language model.🌠'
                              '<a href="https://github.com/Azure99/BlossomLM">GitHub</a>',
                  theme="soft",