tt

app.py

@@ -1,50 +1,46 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-
-# Load the TinyLlama model and tokenizer
-model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
-
-# Initialize the text generation pipeline
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0)
-
-
-def generate(sysA, sysB, wa, wb, user_input):
-    # Construct the system prompts with weights
-    prompt_a = f"System A: {sysA}\n" * int(wa)
-    prompt_b = f"System B: {sysB}\n" * int(wb)
-
-    # Combine prompts and user input
-    full_prompt = prompt_a + prompt_b + f"User: {user_input}\nAssistant:"
-
-    # Generate the response using the model
-    response = generator(full_prompt, max_length=512, num_return_sequences=1)[0]['generated_text']
-
-    return response
-
-
-with gr.Blocks() as demo:
-    gr.Markdown("# Multi-System Prompt Chat Demo")
-
-    with gr.Row():
-        sysA = gr.Textbox(label="System Prompt A", value="You are assistant A.", lines=2)
-        sysB = gr.Textbox(label="System Prompt B", value="You are assistant B.", lines=2)
-
-    with gr.Row():
-        wa = gr.Slider(-5.0, 5.0, value=1.0, step=0.1, label="Weight wA")
-        wb = gr.Slider(-5.0, 5.0, value=1.0, step=0.1, label="Weight wB")
-
-    user_input = gr.Textbox(label="User Message", lines=2)
-    output = gr.Textbox(label="Model Response", lines=10)
-
-    submit_btn = gr.Button("Send")
-
-    submit_btn.click(
-        fn=generate,
-        inputs=[sysA, sysB, wa, wb, user_input],
-        outputs=output
-    )
-
+import torch
+
+# Load tiny model from Hugging Face
+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 if torch.cuda.is_available() else torch.float32
+)
+
+# Use text-generation pipeline (without `device=0`)
+generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+
+# Function to blend two prompts with weights (wa and wb)
+def blend_and_generate(prompt_a, prompt_b, wa, wb):
+    # Normalize weights even if negative
+    total = abs(wa) + abs(wb)
+    if total == 0:
+        return "Error: Both weights cannot be zero."
+    norm_wa = wa / total
+    norm_wb = wb / total
+
+    # Create blended prompt
+    blended_prompt = f"{norm_wa:.2f} * ({prompt_a}) + {norm_wb:.2f} * ({prompt_b})"
+    generated = generator(blended_prompt, max_new_tokens=100, do_sample=True, temperature=0.7)
+    return generated[0]["generated_text"]
+
+# Gradio UI
+demo = gr.Interface(
+    fn=blend_and_generate,
+    inputs=[
+        gr.Textbox(label="Prompt A"),
+        gr.Textbox(label="Prompt B"),
+        gr.Slider(minimum=-5, maximum=5, step=0.1, label="Weight A (wa)"),
+        gr.Slider(minimum=-5, maximum=5, step=0.1, label="Weight B (wb)"),
+    ],
+    outputs=gr.Textbox(label="Generated Output"),
+    title="Tiny Prompt Blender (TinyLlama-1.1B)",
+    description="Enter two prompts and blend them using wa and wb (can be negative).",
+)
+
+# Launch app
 if __name__ == "__main__":
     demo.launch()