Update app.py

app.py

@@ -5,28 +5,45 @@ import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
-MAX_NEW_TOKENS = 2048
-MODEL_NAME = "Azure99/Blossom-V6-7B"
+MAX_NEW_TOKENS = 8192
+MODEL_NAME = "Azure99/Blossom-V6.1-8B"
 
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.bfloat16, device_map="auto")
+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):
     conversation = []
     for user, assistant in history:
-        conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
+        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)
+    return tokenizer.apply_chat_template(conversation, return_tensors="pt").to(
+        model.device
+    )
 
 
-@spaces.GPU
+@spaces.GPU(duration=120)
 def chat(inst, history, temperature, top_p, repetition_penalty):
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    streamer = TextIteratorStreamer(
+        tokenizer, skip_prompt=True, skip_special_tokens=True
+    )
     input_ids = get_input_ids(inst, history)
-    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)
+    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,
+    )
 
     Thread(target=model.generate, kwargs=generation_kwargs).start()
 
@@ -63,23 +80,30 @@ additional_inputs = [
         step=0.01,
         interactive=True,
         info="Repetition Penalty: Controls how much repetition is penalized.",
-    )
+    ),
 ]
 
-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-V6-7B Demo",
-                 description='Hello, I am Blossom, an open source conversational large language model.🌠'
-                             '<a href="https://github.com/Azure99/BlossomLM">GitHub</a>',
-                 theme="soft",
-                 examples=[["Hello"], ["What is MBTI"], ["用Python实现二分查找"],
-                           ["为switch写一篇小红书种草文案，带上emoji"]],
-                 cache_examples=False,
-                 additional_inputs=additional_inputs,
-                 additional_inputs_accordion=gr.Accordion(label="Config", open=True),
-                 clear_btn="🗑️Clear",
-                 undo_btn="↩️Undo",
-                 retry_btn="🔄Retry",
-                 submit_btn="➡️Submit",
-                 ).queue().launch()
+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-V6.1-8B Demo",
+    description="Hello, I am Blossom, an open source conversational large language model.🌠"
+    '<a href="https://github.com/Azure99/BlossomLM">GitHub</a>',
+    theme="soft",
+    examples=[
+        ["Hello"],
+        ["What is MBTI"],
+        ["用Python实现二分查找"],
+        ["为switch写一篇小红书种草文案，带上emoji"],
+    ],
+    cache_examples=False,
+    additional_inputs=additional_inputs,
+    additional_inputs_accordion=gr.Accordion(label="Config", open=True),
+    clear_btn="🗑️Clear",
+    undo_btn="↩️Undo",
+    retry_btn="🔄Retry",
+    submit_btn="➡️Submit",
+).queue().launch()