tt

app.py

@@ -9,57 +9,57 @@ MODEL_ID = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float16, device_map="auto")
 
-def chat_blend(systemA, systemB, wA, wB, user_message, max_new_tokens=50, temperature=1.0, top_p=0.95):
-    # Build messages
+def chat_blend_stream(systemA, systemB, wA, wB, user_message, max_new_tokens=50, temperature=1.0, top_p=0.95):
     promptA = systemA + "\nUser: " + user_message + "\nAssistant:"
     promptB = systemB + "\nUser: " + user_message + "\nAssistant:"
+
     input_ids_A = tokenizer(promptA, return_tensors="pt").input_ids.to(model.device)
     input_ids_B = tokenizer(promptB, return_tensors="pt").input_ids.to(model.device)
-    output_ids = input_ids_A  # share history; or keep separate
+    output_ids = input_ids_A.clone()
+
+    response_text = ""
 
     for _ in range(max_new_tokens):
-        # forward last token
         logitsA = model(input_ids=input_ids_A).logits[:, -1, :]
         logitsB = model(input_ids=input_ids_B).logits[:, -1, :]
+
         probsA = F.softmax(logitsA / temperature, dim=-1)
         probsB = F.softmax(logitsB / temperature, dim=-1)
-        blended = wA * probsA + wB * probsB
-        # apply top_p
-        sorted_probs, sorted_idx = torch.sort(blended, descending=True)
+
+        blended_probs = wA * probsA + wB * probsB
+        sorted_probs, sorted_idx = torch.sort(blended_probs, descending=True)
         cum = torch.cumsum(sorted_probs, dim=-1)
-        mask = cum > top_p
-        sorted_probs[mask] = 0
-        sorted_probs /= sorted_probs.sum(dim=-1, keepdim=True)
-        # repeat mapping back
-        new_id = sorted_idx.gather(-1, torch.multinomial(sorted_probs, num_samples=1))
-        output_ids = torch.cat([output_ids, new_id], dim=-1)
-        # append to each history
-        input_ids_A = torch.cat([input_ids_A, new_id], dim=-1)
-        input_ids_B = torch.cat([input_ids_B, new_id], dim=-1)
-        # stop on EOS
+        sorted_probs[cum > top_p] = 0
+        sorted_probs /= sorted_probs.sum()
+
+        new_id = sorted_idx[0, torch.multinomial(sorted_probs[0], 1)]
+        output_ids = torch.cat([output_ids, new_id.unsqueeze(0).unsqueeze(0)], dim=1)
+        input_ids_A = torch.cat([input_ids_A, new_id.unsqueeze(0).unsqueeze(0)], dim=1)
+        input_ids_B = torch.cat([input_ids_B, new_id.unsqueeze(0).unsqueeze(0)], dim=1)
+
+        token_text = tokenizer.decode(new_id)
+        response_text += token_text
+        yield response_text
+
         if new_id.item() == tokenizer.eos_token_id:
             break
 
-    decoded = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    # strip system + user prompts
-    return decoded.split("Assistant:")[-1].strip()
-
-iface = gr.Interface(
-    fn=chat_blend,
-    inputs=[
+demo = gr.ChatInterface(
+    fn=chat_blend_stream,
+    additional_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="wA", minimum=-2.0, maximum=2.0, step=0.1, value=1.0),
         gr.Slider(label="wB", minimum=-2.0, maximum=2.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=50),
         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.95),
+        gr.Slider(label="Top-p", minimum=0.1, maximum=1.0, step=0.05, value=0.95),
     ],
-    outputs="text",
-    title="Blended‑LLM Chat (TinyLlama)",
-    description="Uses two system prompts and blends their token distributions using wA*p1 + wB*p2."
+    title="Blended TinyLlama",
+    description="Two-system prompts, blended logits with negative/positive weights."
 )
 
+
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch()
+