Update app.py

app.py

@@ -1,278 +1,169 @@
-import torch
-import torch.nn.functional as F
 import gradio as gr
-from typing import List, Tuple
-from transformers import (
-    AutoTokenizer,
-    AutoModelForSeq2SeqLM,
-)
-
-
-# -----------------------------
-# Models
-# -----------------------------
-# Seq2Seq for generation (use selected FLAN-T5) and summarization (DistilBART)
-GEN_MODEL_NAME = "google/flan-t5-small"
-gen_tokenizer = AutoTokenizer.from_pretrained(GEN_MODEL_NAME)
-gen_model = AutoModelForSeq2SeqLM.from_pretrained(GEN_MODEL_NAME)
-
-SUM_MODEL_NAME = "sshleifer/distilbart-cnn-12-6"
-sum_tokenizer = AutoTokenizer.from_pretrained(SUM_MODEL_NAME)
-sum_model = AutoModelForSeq2SeqLM.from_pretrained(SUM_MODEL_NAME)
-
-
-# -----------------------------
-# Helpers
-# -----------------------------
-def _format_token(token_text: str) -> str:
-    """Format a decoded token for table display."""
-    if token_text.strip() == "":
-        # visualize whitespace-only tokens
-        shown = token_text.replace(" ", "␠") or "(space)"
-    else:
-        shown = token_text
-    return f"<code class='tok'>{shown}</code>"
-
-
-def _alternatives_table(
-    chosen_token_ids: List[int],
-    scores: List[torch.Tensor],  # list of [batch, vocab] logits for each step
-    tokenizer,
-    top_k: int = 5,
-    special_ok: bool = False,
-) -> str:
-    """Build an HTML table listing alternatives per generated step.
-
-    chosen_token_ids should align 1:1 with scores.
-    """
-    rows: List[str] = ["<table class='prediction-table'>",
-                       "<tr><th>Generated Token</th><th>Top Alternatives</th></tr>"]
-
-    special_ids = set(tokenizer.all_special_ids or [])
-    for step, (chosen_id, step_scores) in enumerate(zip(chosen_token_ids, scores)):
-        probs = F.softmax(step_scores[0], dim=-1) if step_scores.dim() == 2 else F.softmax(step_scores, dim=-1)
-
-        # Get a surplus then filter out the chosen token and (optionally) specials
-        k_surplus = max(top_k + 10, 20)
-        top_vals, top_idx = torch.topk(probs, k=min(k_surplus, probs.numel()))
-
-        alts: List[Tuple[str, float]] = []
-        for idx, p in zip(top_idx.tolist(), top_vals.tolist()):
-            if idx == chosen_id:
-                continue
-            if not special_ok and idx in special_ids:
-                continue
-            token_text = tokenizer.decode([idx])
-            if token_text == "":
-                continue
-            alts.append((token_text, float(p)))
-            if len(alts) >= top_k:
-                break
-
-        chosen_text = tokenizer.decode([chosen_id])
-        chosen_fmt = _format_token(chosen_text)
-        if alts:
-            alt_fmt = ", ".join(f"{_format_token(t)} <span class='p'>({p*100:.1f}%)</span>" for t, p in alts)
-        else:
-            alt_fmt = "—"
-        rows.append(f"<tr><td class='token-cell'>{chosen_fmt}</td><td class='pred-cell'>{alt_fmt}</td></tr>")
-
-    rows.append("</table>")
-    return "\n".join(rows)
-
-
-# -----------------------------
-# Inference functions
-# -----------------------------
-def generate_with_alternatives(prompt: str, max_tokens: int, temperature: float, top_k: int):
-    if not prompt.strip():
-        return "", ""
-
-    inputs = gen_tokenizer(prompt, return_tensors="pt", truncation=True)
-
-    with torch.no_grad():
-        outputs = gen_model.generate(
-            **inputs,
-            max_new_tokens=int(max_tokens),
-            do_sample=True,
-            temperature=float(temperature),
-            top_p=0.9,
-            repetition_penalty=1.1,
-            return_dict_in_generate=True,
-            output_scores=True,
-        )
-
-    # For encoder-decoder, sequences contain decoder tokens; first is usually decoder_start
-    seq = outputs.sequences[0]
-    text_out = gen_tokenizer.decode(seq, skip_special_tokens=True)
-
-    # Align chosen tokens with scores (skip the decoder start token)
-    chosen_ids = seq[1: 1 + len(outputs.scores)].tolist()
-
-    alt_html = _alternatives_table(
-        chosen_token_ids=chosen_ids,
-        scores=outputs.scores,
-        tokenizer=gen_tokenizer,
-        top_k=int(top_k),
+from transformers import pipeline
+
+# 🌷 Load improved, instruction-tuned models for quality responses
+generator = pipeline("text2text-generation", model="google/flan-t5-small")
+summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")
+
+# 💬 Generate creative or factual text
+def generate_text(prompt, max_tokens, temperature):
+    output = generator(
+        prompt,
+        max_new_tokens=int(max_tokens),
+        temperature=float(temperature),
+        top_p=0.9,
+        repetition_penalty=2.0,
+        do_sample=True,          # 💖 enables randomness and exploration
+        num_return_sequences=1
     )
-
-    alt_wrapper = f"""
-    <div class='response-box fade-in'>
-        <span class='response-label'>Alternative Next Tokens (Generation)</span>
-        {alt_html}
-    </div>
+    return output[0]["generated_text"]
+
+
+
+
+# ✂️ Summarize text
+def summarize_text(text):
+    summary = summarizer(text, max_length=150, min_length=40, do_sample=False)
+    return summary[0]["summary_text"]
+
+# 🌸 Build Gradio Interface
+with gr.Blocks(
+    theme=gr.themes.Soft(primary_hue="pink"),
+    css="""
+        body {
+            background: linear-gradient(135deg, #f3cadb, #f6b8d2, #f7a8c9);
+            font-family: 'Poppins', sans-serif;
+            color: #3b1f2b;
+        }
+
+        h1 {
+            text-align: center;
+            font-size: 2.5em;
+            color: #8e2c4a;
+            text-shadow: 0 0 6px rgba(142,44,74,0.25);
+        }
+
+        p {
+            text-align: center;
+            color: #3b1f2b !important;
+        }
+
+        /* Info box */
+        .helper-box {
+            background: #f8d7e2;
+            border-radius: 18px;
+            padding: 20px;
+            margin: 25px auto;
+            width: 85%;
+            box-shadow: 0 0 12px rgba(122, 41, 64, 0.25);
+            line-height: 1.6;
+            color: #3b1f2b;
+        }
+
+        .helper-box h3 {
+            color: #8e2c4a;
+            text-align: center;
+            font-size: 1.3em;
+        }
+
+        .helper-box b {
+            color: #6e1b3f;
+        }
+
+        .helper-box i, .helper-box span, .helper-box li {
+            color: #3b1f2b !important;
+        }
+
+        /* Labels */
+        label {
+            color: #4a1e33 !important;
+            font-weight: 600;
+        }
+
+        /* Input areas */
+        .gr-textbox textarea {
+            background-color: #f2c8da !important;
+            border-radius: 12px !important;
+            color: #2b0f1a !important;
+        }
+
+        #response-box textarea {
+            min-height: 220px !important;
+            background-color: #f3c6d9 !important;
+            border-radius: 12px !important;
+            padding: 10px !important;
+            font-size: 1em !important;
+            color: #2b0f1a !important;
+            box-shadow: inset 0 0 6px rgba(122, 41, 64, 0.25);
+        }
+
+        /* Buttons */
+        .gr-button {
+            background: #d76c91 !important;
+            color: #2b0f1a !important;
+            font-weight: bold;
+            border-radius: 12px !important;
+            border: 1px solid #8e2c4a !important;
+        }
+
+        .gr-button:hover {
+            transform: scale(1.05);
+            background: #c55b82 !important;
+            box-shadow: 0 0 8px rgba(142,44,74,0.4);
+        }
+
+        /* Tabs */
+        .tab-nav > button {
+            color: #3b1f2b !important;
+            font-weight: 600;
+            background-color: transparent !important;
+        }
+
+        .tab-nav > button.selected {
+            color: #8e2c4a !important;
+            border-bottom: 3px solid #8e2c4a !important;
+        }
+
+        ul li {
+            margin: 6px 0;
+        }
     """
-
-    return text_out, alt_wrapper
-
-
-def summarize_with_alternatives(text: str, max_tokens: int, num_beams: int, top_k: int):
-    if not text.strip():
-        return "", ""
-
-    inputs = sum_tokenizer(text, return_tensors="pt", truncation=True, max_length=1024)
-
-    with torch.no_grad():
-        outputs = sum_model.generate(
-            **inputs,
-            max_new_tokens=int(max_tokens),
-            num_beams=int(num_beams),
-            early_stopping=True,
-            return_dict_in_generate=True,
-            output_scores=True,
-        )
-
-    seq = outputs.sequences[0]
-    summary_text = sum_tokenizer.decode(seq, skip_special_tokens=True)
-
-    # For Seq2Seq, scores length should match tokens after decoder start token
-    # Align chosen token ids with scores (skip the first token which is usually decoder_start)
-    chosen_ids = seq[1: 1 + len(outputs.scores)].tolist()
-
-    alt_html = _alternatives_table(
-        chosen_token_ids=chosen_ids,
-        scores=outputs.scores,
-        tokenizer=sum_tokenizer,
-        top_k=int(top_k),
-    )
-
-    alt_wrapper = f"""
-    <div class='response-box fade-in'>
-        <span class='response-label'>Alternative Next Tokens (Summarization)</span>
-        {alt_html}
-    </div>
-    """
-
-    return summary_text, alt_wrapper
-
-
-# -----------------------------
-# UI
-# -----------------------------
-custom_css = """
-body {
-    background: linear-gradient(135deg, #f3cadb, #f6b8d2, #f7a8c9);
-}
-.gradio-container { background: transparent !important; }
-
-textarea, input {
-    background-color: #f2c8da !important;
-    color: #2b0f1a !important;
-    border-radius: 12px !important;
-    border: 1px solid #8e2c4a !important;
-    font-size: 16px !important;
-}
-button {
-    background: #d76c91 !important;
-    color: #2b0f1a !important;
-    font-weight: bold !important;
-    border-radius: 12px !important;
-    border: 1px solid #8e2c4a !important;
-}
-button:hover {
-    transform: scale(1.03);
-    background: #c55b82 !important;
-}
-
-.response-box {
-    background-color: rgba(255, 230, 235, 0.9);
-    border: 2px solid #ffb6c1;
-    border-radius: 15px;
-    padding: 16px;
-    color: #4b2b30;
-    font-size: 15px;
-    line-height: 1.6;
-    margin-top: 10px;
-    box-shadow: 3px 4px 8px rgba(255, 182, 193, 0.3);
-    white-space: normal;
-    width: 100%;
-    opacity: 0;
-    transition: opacity 0.8s ease-in-out;
-}
-.fade-in { opacity: 1 !important; }
-.response-label {
-    font-weight: bold;
-    color: #d36b83;
-    font-size: 18px;
-    font-family: "Poppins", sans-serif;
-    margin-bottom: 8px;
-    display: block;
-    border-bottom: 1.5px dashed #ffc9d6;
-    padding-bottom: 5px;
-    text-align: center;
-}
-
-.prediction-table { width: 100%; border-collapse: collapse; margin-top: 6px; }
-.prediction-table th {
-    background-color: #ffe6eb;
-    color: #d36b83;
-    text-align: center;
-    font-family: "Poppins", sans-serif;
-    padding: 6px;
-    border-bottom: 2px dashed #ffc9d6;
-}
-.prediction-table td { padding: 6px 10px; color: #4b2b30; border-bottom: 1px solid #ffd3de; vertical-align: top; }
-.token-cell { font-weight: bold; text-align: right; width: 25%; color: #c45c77; padding-right: 10px; }
-.pred-cell { text-align: left; width: 75%; }
-.tok { font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, "Liberation Mono", "Courier New", monospace; }
-.p { color: #997; }
-"""
-
-
-with gr.Blocks(css=custom_css) as demo:
-    gr.Markdown("<h1>🌸 App 3: Generate & Summarize with Alternatives</h1>")
-    gr.Markdown("<p>Create text or summaries, and inspect alternative tokens considered at each step.</p>")
-
+) as app:
+
+    # 🌸 Title
+    gr.Markdown("<h1>🌸 AI Assistant</h1>")
+    gr.Markdown("<p>✨ Generate or Summarize Text using Pretrained Hugging Face Models ✨</p>")
+
+    # 💡 Instructions Box
+    gr.HTML("""
+        <div class="helper-box">
+            <h3>💡 How It Works</h3>
+            <p style="color:#3b1f2b;">
+                Choose between <b>text generation</b> or <b>text summarization</b> below!<br>
+                This app uses improved open-source models for more natural and reliable text output.
+            </p>
+            <ul>
+                <li><b>Text Generation:</b> <span>Type a creative prompt, like <i>“Write a story about a butterfly in Texas.”</i></span></li>
+                <li><b>Text Summarization:</b> <span>Paste a paragraph or article to get a clear, shorter summary.</span></li>
+            </ul>
+            <p style="color:#3b1f2b;">🎀 <b>Max Tokens:</b> Longer responses use more tokens.</p>
+            <p style="color:#3b1f2b;">🌡️ <b>Temperature:</b> Controls creativity — lower = logical, higher = imaginative.</p>
+        </div>
+    """)
+
+    # 💬 Text Generation Tab
     with gr.Tab("💬 Text Generation"):
-        gen_prompt = gr.Textbox(label="Enter your prompt:")
-        with gr.Row():
-            gen_max_tokens = gr.Slider(10, 300, value=120, step=10, label="Max New Tokens")
-            gen_temperature = gr.Slider(0.1, 1.5, value=0.8, step=0.1, label="Temperature")
-            gen_top_k = gr.Slider(2, 10, value=5, step=1, label="Top-k Alternatives")
-        with gr.Row():
-            gen_output = gr.Textbox(label="Generated Output", lines=8)
-        gen_alts = gr.HTML(label="Alternatives", value="")
-        gr.Button("✨ Generate").click(
-            generate_with_alternatives,
-            inputs=[gen_prompt, gen_max_tokens, gen_temperature, gen_top_k],
-            outputs=[gen_output, gen_alts],
-        )
+        prompt = gr.Textbox(label="💖 Enter your prompt:")
+        max_tokens = gr.Slider(50, 500, value=150, step=10, label="🎀 Max Tokens")
+        temperature = gr.Slider(0.1, 1.0, value=0.7, step=0.1, label="🌡️ Temperature")
+        generation_output = gr.Textbox(label="💬 Generated Response", lines=10, elem_id="response-box")
+        gr.Button("✨ Generate ✨").click(generate_text, [prompt, max_tokens, temperature], generation_output)
 
+    # 🧁 Text Summarization Tab
     with gr.Tab("🧁 Text Summarization"):
-        sum_text = gr.Textbox(label="Paste your text:", lines=8)
-        with gr.Row():
-            sum_max_tokens = gr.Slider(20, 200, value=120, step=10, label="Max New Tokens")
-            sum_beams = gr.Slider(1, 6, value=4, step=1, label="Beams")
-            sum_top_k = gr.Slider(2, 10, value=5, step=1, label="Top-k Alternatives")
-        with gr.Row():
-            sum_output = gr.Textbox(label="Summary", lines=8)
-        sum_alts = gr.HTML(label="Alternatives", value="")
-        gr.Button("🎀 Summarize").click(
-            summarize_with_alternatives,
-            inputs=[sum_text, sum_max_tokens, sum_beams, sum_top_k],
-            outputs=[sum_output, sum_alts],
-        )
-
+        input_text = gr.Textbox(label="📝 Paste your text:", lines=10)
+        summary_output = gr.Textbox(label="🍬 Summary", lines=10, elem_id="response-box")
+        gr.Button("🎀 Summarize 🎀").click(summarize_text, [input_text], summary_output)
 
-if __name__ == "__main__":
-    demo.launch()
+# 🚀 Launch
+app.launch()