tt

app.py

@@ -1,16 +1,26 @@
 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):
-    # Simple blending: concatenate weighted system prompts with user input
-    # For example, we just repeat sysA and sysB according to weights (rounded)
-    prompt_a = (sysA + " ") * max(0, int(round(wa)))
-    prompt_b = (sysB + " ") * max(0, int(round(wb)))
+    # 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:"
 
-    blended_prompt = prompt_a + prompt_b + "\nUser: " + user_input
+    # Generate the response using the model
+    response = generator(full_prompt, max_length=512, num_return_sequences=1)[0]['generated_text']
 
-    # Pretend model response is just echoing blended prompt
-    response = f"Blended prompt sent to model:\n{blended_prompt}"
     return response