tt

app.py

@@ -17,37 +17,41 @@ model_b = AutoModelForCausalLM.from_pretrained(model_name_b).to(device)
 model_a.eval()
 model_b.eval()
 
-def blend_generate(prompt, wa, wb):
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+def blend_generate(system_prompt_a, system_prompt_b, user_prompt, wa, wb):
+    # Combine system prompt A + user prompt for model A
+    prompt_a = system_prompt_a + user_prompt
+    # Combine system prompt B + user prompt for model B
+    prompt_b = system_prompt_b + user_prompt
+
+    input_ids_a = tokenizer(prompt_a, return_tensors="pt").input_ids.to(device)
+    input_ids_b = tokenizer(prompt_b, return_tensors="pt").input_ids.to(device)
 
     with torch.no_grad():
-        output_a = model_a(input_ids)
-        output_b = model_b(input_ids)
+        output_a = model_a(input_ids_a)
+        output_b = model_b(input_ids_b)
 
     logits_a = output_a.logits[:, -1, :]
     logits_b = output_b.logits[:, -1, :]
 
-    # Weighted sum of raw logits (before softmax)
     blended_logits = wa * logits_a + wb * logits_b
 
-    # Softmax to get probabilities
     probs = torch.softmax(blended_logits, dim=-1)
-
-    # Sample one token from the blended distribution
     token = torch.multinomial(probs, 1)
     next_token_id = token.item()
-
     next_token = tokenizer.decode([next_token_id])
-    return prompt + next_token
 
-# Gradio UI
+    # For simplicity, just return user prompt + next token (you can customize)
+    return user_prompt + next_token
+
 with gr.Blocks() as demo:
-    prompt_input = gr.Textbox(label="Prompt", lines=2)
-    weight_a = gr.Slider(0, 1, value=0.5, label="Weight model A")
-    weight_b = gr.Slider(0, 1, value=0.5, label="Weight model B")
+    system_prompt_a = gr.Textbox(label="System Prompt A", value="You are a helpful assistant. ")
+    system_prompt_b = gr.Textbox(label="System Prompt B", value="You are a sarcastic assistant. ")
+    user_prompt = gr.Textbox(label="User Prompt")
+    weight_a = gr.Slider(0, 1, value=0.5, label="Weight Model A")
+    weight_b = gr.Slider(0, 1, value=0.5, label="Weight Model B")
     output_text = gr.Textbox(label="Output")
 
     btn = gr.Button("Generate")
-    btn.click(blend_generate, inputs=[prompt_input, weight_a, weight_b], outputs=output_text)
+    btn.click(blend_generate, inputs=[system_prompt_a, system_prompt_b, user_prompt, weight_a, weight_b], outputs=output_text)
 
 demo.launch()