Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from transformers import pipeline | |
| # =============================== | |
| # Load open-source text generation models | |
| # =============================== | |
| models = { | |
| "DistilGPT-2": "distilgpt2", | |
| "GPT2 (Small)": "gpt2", | |
| "DialoGPT-small": "microsoft/DialoGPT-small", | |
| "OPT-350M": "facebook/opt-350m", | |
| "Bloom-560M": "bigscience/bloom-560m", | |
| "GPT-Neo-125M": "EleutherAI/gpt-neo-125M", | |
| "Falcon-RW-1B": "tiiuae/falcon-rw-1b", | |
| "Flan-T5-Small": "google/flan-t5-small", | |
| "Flan-T5-Base": "google/flan-t5-base", | |
| "Phi-2": "microsoft/phi-2" | |
| } | |
| generators = {name: pipeline("text-generation", model=mdl) | |
| if "flan" not in mdl.lower() and "bart" not in mdl.lower() | |
| else pipeline("text2text-generation", model=mdl) | |
| for name, mdl in models.items()} | |
| # Summarizer model | |
| summarizer = pipeline("text2text-generation", model="google/flan-t5-base") | |
| # =============================== | |
| # Function to query all models | |
| # =============================== | |
| def compare_models(user_input, max_new_tokens=100, temperature=0.7): | |
| raw_outputs = {} | |
| clean_outputs = {} | |
| for name, generator in generators.items(): | |
| try: | |
| if "text-generation" in generator.task: | |
| output = generator( | |
| user_input, | |
| max_new_tokens=max_new_tokens, | |
| temperature=temperature | |
| )[0]["generated_text"] | |
| else: # Flan models etc | |
| output = generator(user_input, max_new_tokens=max_new_tokens)[0]["generated_text"] | |
| raw_outputs[name] = output | |
| # Summarize the answer to improve clarity | |
| summary = summarizer("Summarize this: " + output, max_new_tokens=60)[0]["generated_text"] | |
| clean_outputs[name] = summary | |
| except Exception as e: | |
| raw_outputs[name] = f"⚠️ Error: {str(e)}" | |
| clean_outputs[name] = "N/A" | |
| return [raw_outputs[m] for m in models.keys()], [clean_outputs[m] for m in models.keys()] | |
| # =============================== | |
| # Gradio UI | |
| # =============================== | |
| with gr.Blocks(css="style.css") as demo: | |
| gr.Markdown("## 🤖 Open-Source Model Comparator\n" | |
| "Compare outputs from multiple open-source LLMs side by side.\n" | |
| "Includes a raw output and a cleaned summary (via Flan-T5).") | |
| with gr.Row(): | |
| user_input = gr.Textbox(label="Your prompt", placeholder="Try: 'Explain quantum computing in simple terms'", lines=2) | |
| generate_btn = gr.Button("Generate", variant="primary") | |
| with gr.Row(): | |
| max_tokens = gr.Slider(20, 200, value=100, step=10, label="Max new tokens") | |
| temp = gr.Slider(0.1, 1.0, value=0.7, step=0.1, label="Creativity (temperature)") | |
| gr.Markdown("### 🔎 Raw Outputs") | |
| with gr.Row(): | |
| raw_boxes = [gr.Textbox(label=name, elem_classes="output-box", interactive=False) for name in models.keys()] | |
| gr.Markdown("### ✨ Cleaned Summaries (Flan-T5)") | |
| with gr.Row(): | |
| clean_boxes = [gr.Textbox(label=f"{name} (Summary)", elem_classes="output-box", interactive=False) for name in models.keys()] | |
| examples = [ | |
| ["Explain quantum computing in simple terms."], | |
| ["Write a haiku about autumn leaves."], | |
| ["What are the pros and cons of nuclear energy?"], | |
| ["Describe a futuristic city in the year 2200."], | |
| ["Write a funny short story about a robot learning to cook."] | |
| ] | |
| gr.Examples(examples=examples, inputs=[user_input]) | |
| generate_btn.click(compare_models, inputs=[user_input, max_tokens, temp], outputs=raw_boxes + clean_boxes) | |
| user_input.submit(compare_models, inputs=[user_input, max_tokens, temp], outputs=raw_boxes + clean_boxes) | |
| if __name__ == "__main__": | |
| demo.launch() | |