tt

app.py

@@ -6,8 +6,9 @@ import gradio as gr
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Load two small models and their tokenizer (you can replace these with your models)
-model_name_a = "distilgpt2"
-model_name_b = "sshleifer/tiny-gpt2"  # very small GPT2 variant for demo
+MODEL_NAME = "arnir0/Tiny-LLM"
+model_name_a = MODEL_NAME
+model_name_b = MODEL_NAME
 
 tokenizer = AutoTokenizer.from_pretrained(model_name_a)
 
@@ -17,31 +18,39 @@ model_b = AutoModelForCausalLM.from_pretrained(model_name_b).to(device)
 model_a.eval()
 model_b.eval()
 
-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
+def blend_generate(system_prompt_a, system_prompt_b, user_prompt, wa, wb, max_length=50):
+    generated_text = user_prompt
+    device = next(model_a.parameters()).device  # infer device from model
 
-    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)
+    for _ in range(max_length):
+        # Prepare prompts for each model: system prompt + generated text so far
+        prompt_a = system_prompt_a + generated_text
+        prompt_b = system_prompt_b + generated_text
 
-    with torch.no_grad():
-        output_a = model_a(input_ids_a)
-        output_b = model_b(input_ids_b)
+        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)
 
-    logits_a = output_a.logits[:, -1, :]
-    logits_b = output_b.logits[:, -1, :]
+        with torch.no_grad():
+            output_a = model_a(input_ids_a)
+            output_b = model_b(input_ids_b)
 
-    blended_logits = wa * logits_a + wb * logits_b
+        logits_a = output_a.logits[:, -1, :]
+        logits_b = output_b.logits[:, -1, :]
 
-    probs = torch.softmax(blended_logits, dim=-1)
-    token = torch.multinomial(probs, 1)
-    next_token_id = token.item()
-    next_token = tokenizer.decode([next_token_id])
+        blended_logits = wa * logits_a + wb * logits_b
 
-    # For simplicity, just return user prompt + next token (you can customize)
-    return user_prompt + next_token
+        probs = torch.softmax(blended_logits, dim=-1)
+        token = torch.multinomial(probs, 1)
+        next_token_id = token.item()
+
+        # Stop if end-of-sequence token generated (adjust based on your tokenizer)
+        if next_token_id == tokenizer.eos_token_id:
+            break
+
+        next_token = tokenizer.decode([next_token_id])
+        generated_text += next_token
+
+    return generated_text
 
 with gr.Blocks() as demo:
     system_prompt_a = gr.Textbox(label="System Prompt A", value="You are a helpful assistant. ")