tt

app.py

@@ -11,7 +11,7 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 model.eval()
 
-def generate_stream(sysA, sysB, wA, wB, user_message, max_new_tokens=100, temperature=1.0, top_p=0.9):
+def blend_generate(sysA, sysB, wA, wB, user_message, max_new_tokens, temperature, top_p):
     promptA = f"<|system|>{sysA}\n<|user|>{user_message}\n<|assistant|>"
     promptB = f"<|system|>{sysB}\n<|user|>{user_message}\n<|assistant|>"
 
@@ -20,7 +20,6 @@ def generate_stream(sysA, sysB, wA, wB, user_message, max_new_tokens=100, temper
 
     outA, outB = idsA.clone(), idsB.clone()
     response = ""
-    yield response  # start stream
 
     for _ in range(max_new_tokens):
         with torch.no_grad():
@@ -28,7 +27,7 @@ def generate_stream(sysA, sysB, wA, wB, user_message, max_new_tokens=100, temper
             logitsB = model(input_ids=outB).logits[:, -1, :]
 
         blended = wA * logitsA + wB * logitsB
-        blended = blended / (temperature if temperature > 0 else 1.0)
+        blended = blended / temperature
 
         probs = F.softmax(blended, dim=-1)
         sorted_probs, sorted_idx = torch.sort(probs, descending=True)
@@ -47,20 +46,26 @@ def generate_stream(sysA, sysB, wA, wB, user_message, max_new_tokens=100, temper
         if token.item() == tokenizer.eos_token_id:
             break
 
-demo = gr.ChatInterface(
-    fn=generate_stream,
-    inputs=[
-        gr.Textbox(label="System Prompt A", value="You are assistant A."),
-        gr.Textbox(label="System Prompt B", value="You are assistant B."),
-        gr.Slider(label="Weight wA", minimum=-5.0, maximum=5.0, step=0.1, value=1.0),
-        gr.Slider(label="Weight wB", minimum=-5.0, maximum=5.0, step=0.1, value=1.0),
-        gr.Textbox(label="User Message"),
-        gr.Slider(label="Max New Tokens", minimum=1, maximum=200, step=1, value=100),
-        gr.Slider(label="Temperature", minimum=0.1, maximum=2.0, step=0.1, value=1.0),
-        gr.Slider(label="Top-p", minimum=0.1, maximum=1.0, step=0.05, value=0.9),
-    ],
-    title="Streaming Blended TinyLlama Chat"
-)
+with gr.Blocks() as demo:
+    gr.Markdown("## Blended Prompt Chat (TinyLlama)")
+    sysA = gr.Textbox(label="System Prompt A", value="You are assistant A.")
+    sysB = gr.Textbox(label="System Prompt B", value="You are assistant B.")
+    wA = gr.Slider(-5, 5, value=1.0, step=0.1, label="Weight A")
+    wB = gr.Slider(-5, 5, value=1.0, step=0.1, label="Weight B")
+    user_msg = gr.Textbox(label="User Message")
+    temp = gr.Slider(0.1, 2.0, value=1.0, step=0.1, label="Temperature")
+    top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
+    max_tokens = gr.Slider(1, 200, value=100, step=1, label="Max New Tokens")
+    output = gr.Textbox(label="Response")
+
+    btn = gr.Button("Generate")
+    btn.click(
+        blend_generate,
+        [sysA, sysB, wA, wB, user_msg, max_tokens, temp, top_p],
+        output,
+        show_progress=True,
+        stream=True
+    )
 
 if __name__ == "__main__":
     demo.launch()

requirements.txt

@@ -1,4 +1,3 @@
-transformers>=4.31
-torch
-gradio
-accelerate
+gradio>=3.50.0
+transformers>=4.40.0
+torch>=2.2.0