import gradio as gr
import torch
from transformers import pipeline

MODEL_ID = "beta3/gemma3_1b_title_generator"

# Load pipeline
pipe = pipeline(
    "text-generation",
    model=MODEL_ID,
    dtype=torch.bfloat16,
    device_map="auto"
)

def generate_title(abstract, temperature=0.7, top_p=0.9, max_tokens=32):
    if not abstract or not abstract.strip():
        return "Please provide a research abstract."

    prompt = f"""<bos><start_of_turn>user
Generate a concise academic title for the following abstract:
{abstract}
<end_of_turn>
<start_of_turn>model
"""

    output = pipe(
        prompt,
        max_new_tokens=max_tokens,
        do_sample=True,
        temperature=temperature,
        top_p=top_p,
        return_full_text=False
    )

    return output[0].get("generated_text", "").strip()


with gr.Blocks(title="Academic Title Generator · Gemma 3") as demo:

    # ===== Header =====
    gr.Markdown(
        """
        # Academic Title Generator

        This demo generates concise academic paper titles from research abstracts.
        The model is based on **Gemma 3 (1B)** and fine-tuned specifically for
        academic title generation using **LoRA**.
        """
    )

    gr.Markdown("---")

    # ===== Main layout =====
    with gr.Row(equal_height=True):
        with gr.Column(scale=3):
            abstract_input = gr.Textbox(
                lines=10,
                label="Research Abstract",
                value=(
                    "Transformer-based architectures have demonstrated strong performance "
                    "in tasks involving reasoning, scientific understanding, and text generation. "
                    "Producing concise academic titles from long abstracts, however, remains a "
                    "non-trivial task."
                ),
                placeholder="Paste your research abstract here..."
            )

            generate_button = gr.Button(
                "Generate title",
                variant="primary"
            )

        with gr.Column(scale=2):
            output_title = gr.Textbox(
                label="Generated title",
                placeholder="The generated title will appear here.",
                lines=4
            )

            gr.Markdown(
                """
                **Usage notes**

                - The model produces a single concise academic title per request.
                - Outputs may vary slightly between runs due to probabilistic sampling.
                - The model is optimized for formal academic and scientific writing.
                - Best suited for research papers, preprints, and technical reports.
                """
            )

    # ===== Advanced settings =====
    if hasattr(gr, "Accordion"):
        with gr.Accordion("Advanced generation settings", open=False):
            with gr.Row():
                temperature = gr.Slider(
                    0.1, 1.5,
                    value=0.7,
                    step=0.05,
                    label="Temperature",
                    info="Higher values increase variability in the generated title."
                )
                top_p = gr.Slider(
                    0.1, 1.0,
                    value=0.9,
                    step=0.05,
                    label="Top-p"
                )
                max_tokens = gr.Slider(
                    8, 64,
                    value=32,
                    step=1,
                    label="Maximum new tokens"
                )
    else:
        gr.Markdown("**Advanced generation settings**")
        with gr.Row():
            temperature = gr.Slider(0.1, 1.5, value=0.7, step=0.05, label="Temperature")
            top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
            max_tokens = gr.Slider(8, 64, value=32, step=1, label="Maximum new tokens")

    # ===== Footer =====
    gr.Markdown(
        """
        ---
        **Model**: Gemma 3 (1B)  
        **Fine-tuning**: LoRA (Unsloth)  
        **Task**: Academic title generation  

        This is a research demo intended for exploratory and experimental use.
        """
    )

    generate_button.click(
        generate_title,
        inputs=[abstract_input, temperature, top_p, max_tokens],
        outputs=output_title
    )

demo.launch(
    server_name="0.0.0.0",
    server_port=7860,
)