Spaces:
Sleeping
Sleeping
| import os | |
| import gradio as gr | |
| import torch | |
| from transformers import pipeline, BitsAndBytesConfig | |
| # 1. Setup Auth | |
| token = os.environ.get("HF_TOKEN") | |
| # 2. Configuration for 4-bit Quantization (Saves RAM) | |
| bnb_config = BitsAndBytesConfig( | |
| load_in_4bit=True, | |
| bnb_4bit_compute_dtype=torch.bfloat16 | |
| ) | |
| print("Loading 3 models into RAM... please wait about 5 minutes.") | |
| # Load Gemma 3 1B (Uses ~1GB RAM) | |
| pipe_1b = pipeline("text-generation", model="google/gemma-3-1b-it", | |
| device_map="cpu", token=token, model_kwargs={"quantization_config": bnb_config}) | |
| # Load Gemma 3 4B (Uses ~3.5GB RAM in 4-bit) | |
| pipe_4b = pipeline("text-generation", model="google/gemma-3-4b-it", | |
| device_map="cpu", token=token, model_kwargs={"quantization_config": bnb_config}) | |
| # Load Phi-4 Mini (Uses ~2.5GB RAM in 4-bit) | |
| pipe_phi = pipeline("text-generation", model="microsoft/Phi-4-mini-instruct", | |
| device_map="cpu", token=token, model_kwargs={"quantization_config": bnb_config}) | |
| def battle(message): | |
| results = [] | |
| # Format message for Instruct models | |
| msgs = [{"role": "user", "content": message}] | |
| # Generate from all three | |
| out_1b = pipe_1b(msgs, max_new_tokens=128)[0]['generated_text'][-1]['content'] | |
| out_4b = pipe_4b(msgs, max_new_tokens=128)[0]['generated_text'][-1]['content'] | |
| out_phi = pipe_phi(msgs, max_new_tokens=128)[0]['generated_text'][-1]['content'] | |
| return out_1b, out_4b, out_phi | |
| # 3. UI Layout | |
| with gr.Blocks(theme=gr.themes.Soft()) as demo: | |
| gr.Markdown("# ⚔️ Triple Model Local Arena") | |
| gr.Markdown("Running locally on 16GB CPU RAM via 4-bit Quantization.") | |
| with gr.Row(): | |
| txt = gr.Textbox(label="Prompt", placeholder="Compare their logic...", scale=4) | |
| btn = gr.Button("Battle!", variant="primary", scale=1) | |
| with gr.Row(): | |
| with gr.Column(): | |
| gr.Markdown("### ⚡ Gemma 3 1B") | |
| res_1 = gr.Markdown() | |
| with gr.Column(): | |
| gr.Markdown("### 🧠 Gemma 3 4B") | |
| res_2 = gr.Markdown() | |
| with gr.Column(): | |
| gr.Markdown("### 🛠️ Phi-4 Mini") | |
| res_3 = gr.Markdown() | |
| btn.click(battle, inputs=txt, outputs=[res_1, res_2, res_3]) | |
| if __name__ == "__main__": | |
| demo.launch(server_name="0.0.0.0", server_port=7860) |