tt

app.py

@@ -1,72 +1,35 @@
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-model_name_a = "mistralai/Mistral-7B-v0.1"
-model_name_b = "mistralai/Mistral-7B-v0.1"  # you can replace this with a second different model or finetuned variant
+MODEL_NAME = "microsoft/Phi-4-mini-instruct"
 
-print("Loading tokenizer...")
-tokenizer = AutoTokenizer.from_pretrained(model_name_a)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 
-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)
+def generate_text(prompt, model, tokenizer, max_length=512, temperature=1, top_k=50, top_p=0.95):
+    inputs = tokenizer.encode(prompt, return_tensors="pt")
 
-model_a.eval()
-model_b.eval()
+    outputs = model.generate(
+        inputs,
+        max_length=max_length,
+        temperature=temperature,
+        top_k=top_k,
+        top_p=top_p,
+        do_sample=True
+    )
 
-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
-
-    for _ in range(max_length):
-        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)
-
-        with torch.no_grad():
-            output_a = model_a(input_ids_a)
-            output_b = model_b(input_ids_b)
-
-        logits_a = output_a.logits[:, -1, :]
-        logits_b = output_b.logits[:, -1, :]
-
-        blended_logits = wa * logits_a + wb * logits_b
-
-        # 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)
-
-        # 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_str = tokenizer.decode([next_token])
-        generated_text += next_token_str
 
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return generated_text
 
+def main():
+    # Define your prompt
+    prompt = "According to all known laws of aviation, there is no way a bee should be able to fly."
 
-with gr.Blocks() as demo:
-    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)
+    generated_text = generate_text(prompt, model, tokenizer)
 
-    btn = gr.Button("Generate")
-    btn.click(blend_generate, inputs=[system_prompt_a, system_prompt_b, user_prompt, weight_a, weight_b], outputs=output_text)
+    print(generated_text)
 
-demo.launch()
+if __name__ == "__main__":
+    main()