tt

app.py

@@ -1,65 +1,70 @@
 import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import torch.nn.functional as F
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-# Model selection
-MODEL_ID = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 
+# Load tiny model (CPU-friendly)
+MODEL_ID = "tiiuae/falcon-rw-1b"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float16, device_map="auto")
-
-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.clone()
-
-    response_text = ""
-
-    for _ in range(max_new_tokens):
-        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_probs = wA * probsA + wB * probsB
-        sorted_probs, sorted_idx = torch.sort(blended_probs, descending=True)
-        cum = torch.cumsum(sorted_probs, dim=-1)
-        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
-
-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.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),
-    ],
-    title="Blended TinyLlama",
-    description="Two-system prompts, blended logits with negative/positive weights."
-)
-
-
-if __name__ == "__main__":
-    demo.launch()
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID).to("cpu")
+
+
+def get_logits(prompt, system_msg):
+    """Run the model and return logits for the next token."""
+    input_text = f"<|system|>{system_msg}\n<|user|>{prompt}<|assistant|>"
+    inputs = tokenizer(input_text, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+    logits = outputs.logits[:, -1, :]  # Only final token logits
+    return logits
+
+
+def blended_generate(prompt, sys1, sys2, wa, wb, temperature=1.0):
+    # Get logits from both system prompts
+    logits1 = get_logits(prompt, sys1)
+    logits2 = get_logits(prompt, sys2)
+
+    # Weighted sum
+    blended_logits = wa * logits1 + wb * logits2
+
+    # Apply temperature
+    blended_logits = blended_logits / temperature
+
+    # Convert to probabilities
+    probs = F.softmax(blended_logits, dim=-1)
+
+    # Sample one token from the distribution
+    token_id = torch.multinomial(probs, num_samples=1)
+    next_token = tokenizer.decode(token_id[0])
+
+    return next_token.strip()
+
+
+# Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("## 🔀 Blended System Prompts using Falcon-RW-1B")
+
+    with gr.Row():
+        prompt = gr.Textbox(label="User Prompt", value="Tell me a joke about computers.")
+
+    with gr.Row():
+        sys1 = gr.Textbox(label="System Prompt A", value="You are a polite assistant.")
+        sys2 = gr.Textbox(label="System Prompt B", value="You are a sarcastic assistant.")
+
+    with gr.Row():
+        wa = gr.Slider(-10, 10, value=1.0, step=0.1, label="Weight A")
+        wb = gr.Slider(-10, 10, value=1.0, step=0.1, label="Weight B")
+
+    with gr.Row():
+        temperature = gr.Slider(0.1, 2.0, value=1.0, step=0.1, label="Temperature")
+
+    output = gr.Textbox(label="Next Token")
+
+    generate_btn = gr.Button("Generate Next Token")
+    generate_btn.click(
+        fn=blended_generate,
+        inputs=[prompt, sys1, sys2, wa, wb, temperature],
+        outputs=output,
+    )
 
+demo.launch()