tt

app.py

@@ -1,31 +1,29 @@
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
 
-# Set device: GPU if available, else CPU
 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 = "arnir0/Tiny-LLM"
-model_name_a = MODEL_NAME
-model_name_b = MODEL_NAME
+model_name_a = "meta-llama/Llama-2-7b-chat-hf"
+model_name_b = "meta-llama/Llama-2-7b-chat-hf"  # you can replace this with a second different model or finetuned variant
 
+print("Loading tokenizer...")
 tokenizer = AutoTokenizer.from_pretrained(model_name_a)
 
-model_a = AutoModelForCausalLM.from_pretrained(model_name_a).to(device)
-model_b = AutoModelForCausalLM.from_pretrained(model_name_b).to(device)
+print("Loading models...")
+model_a = AutoModelForCausalLM.from_pretrained(model_name_a, device_map="auto", torch_dtype=torch.float16)
+model_b = AutoModelForCausalLM.from_pretrained(model_name_b, device_map="auto", torch_dtype=torch.float16)
 
 model_a.eval()
 model_b.eval()
 
-def blend_generate(system_prompt_a, system_prompt_b, user_prompt, wa, wb, max_length=50):
+def blend_generate(system_prompt_a, system_prompt_b, user_prompt, wa, wb, max_length=50, temperature=0.7, top_k=50):
+    device = next(model_a.parameters()).device
     generated_text = user_prompt
-    device = next(model_a.parameters()).device  # infer device from model
 
     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
+        prompt_a = system_prompt_a.strip() + "\n" + generated_text
+        prompt_b = system_prompt_b.strip() + "\n" + generated_text
 
         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)
@@ -39,26 +37,34 @@ def blend_generate(system_prompt_a, system_prompt_b, user_prompt, wa, wb, max_le
 
         blended_logits = wa * logits_a + wb * logits_b
 
-        probs = torch.softmax(blended_logits, dim=-1)
-        token = torch.multinomial(probs, 1)
-        next_token_id = token.item()
+        # Apply top-k filtering
+        top_k_logits, top_k_indices = torch.topk(blended_logits, top_k)
+        filtered_logits = torch.full_like(blended_logits, float('-inf'))
+        filtered_logits.scatter_(1, top_k_indices, top_k_logits)
 
-        # Stop if end-of-sequence token generated (adjust based on your tokenizer)
-        if next_token_id == tokenizer.eos_token_id:
+        # Temperature scaling
+        scaled_logits = filtered_logits / temperature
+
+        probs = torch.softmax(scaled_logits, dim=-1)
+
+        next_token = torch.multinomial(probs, 1).item()
+
+        if next_token == tokenizer.eos_token_id:
             break
 
-        next_token = tokenizer.decode([next_token_id])
-        generated_text += next_token
+        next_token_str = tokenizer.decode([next_token])
+        generated_text += next_token_str
 
     return generated_text
 
+
 with gr.Blocks() as demo:
-    system_prompt_a = gr.Textbox(label="System Prompt A", value="You are a funny assistant. ")
-    system_prompt_b = gr.Textbox(label="System Prompt B", value="You are a angry assistant. ")
-    user_prompt = gr.Textbox(label="User Prompt",value="tell me a story")
-    weight_a = gr.Slider(-2, 2, value=1, label="Weight Model A")
-    weight_b = gr.Slider(-2, 2, value=1, label="Weight Model B")
-    output_text = gr.Textbox(label="Output")
+    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 witty assistant.")
+    user_prompt = gr.Textbox(label="User Prompt", value="Tell me a story about a dragon.")
+    weight_a = gr.Slider(minimum=0, maximum=1, value=0.5, label="Weight Model A")
+    weight_b = gr.Slider(minimum=0, maximum=1, value=0.5, label="Weight Model B")
+    output_text = gr.Textbox(label="Output", lines=10)
 
     btn = gr.Button("Generate")
     btn.click(blend_generate, inputs=[system_prompt_a, system_prompt_b, user_prompt, weight_a, weight_b], outputs=output_text)