Update app.py

app.py

@@ -1,4 +1,3 @@
-import spaces
 from threading import Thread
 
 import torch
@@ -7,13 +6,12 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStream
 import time
 
 model_id = "facebook/opt-6.7b"
-assistant_id = "facebook/opt-125m"
+assistant_id = "facebook/opt-350m"
 
 model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")
 assistant_model = AutoModelForCausalLM.from_pretrained(assistant_id).to(device=model.device, dtype=torch.bfloat16)
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
-@spaces.GPU
 def run_generation(user_text, use_assistant, temperature, max_new_tokens):
     if temperature < 0.1:
         do_sample = False
@@ -31,6 +29,7 @@ def run_generation(user_text, use_assistant, temperature, max_new_tokens):
         assistant_model=assistant_model if use_assistant else None,
         streamer=streamer,
         max_new_tokens=max_new_tokens,
+        min_new_tokens=max_new_tokens,
         do_sample=do_sample,
         top_p=0.95,
         temperature=float(temperature),
@@ -65,7 +64,7 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column(scale=4):
             user_text = gr.Textbox(
-                placeholder="A sequence: one, two, three, ",
+                value="A sequence: one, two, three, ",
                 label="Prompt"
             )
             model_output = gr.Textbox(label="Model output", lines=10, interactive=False)
@@ -88,4 +87,4 @@ with gr.Blocks() as demo:
     user_text.submit(run_generation, generate_inputs, generate_outputs)
     button_submit.click(run_generation, generate_inputs, generate_outputs)
 
-    demo.queue(max_size=32).launch()
+    demo.queue(max_size=16).launch()