src/generation.py

from __future__ import annotations

from typing import Tuple

import streamlit as st
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


@st.cache_resource(show_spinner=True)
def load_model_and_tokenizer(model_id: str) -> Tuple[AutoTokenizer, AutoModelForCausalLM]:
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        device_map="auto",
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    )
    return tokenizer, model


def build_user_content(title: str, introduction: str, conclusion: str) -> str:
    return (
        "You are an expert academic writer. Write a concise and high-quality abstract "
        "(150-250 words) from the provided paper context.\n\n"
        f"Title: {title}\n\n"
        f"Introduction:\n{introduction}\n\n"
        f"Conclusion:\n{conclusion}\n\n"
        "Write the abstract now."
    )


def build_user_content_from_article(article_text: str) -> str:
    return (
        "You are an expert academic writing assistant. Summarize the following article text "
        "into a concise, faithful summary (150-250 words) with clear key findings.\n\n"
        f"Article text:\n{article_text}\n\n"
        "Write the summary now."
    )


def build_prompt_gemma_fallback(user_content: str) -> str:
    return f"""<start_of_turn>user
{user_content}<end_of_turn>
<start_of_turn>model
"""


def build_prompt(tokenizer: AutoTokenizer, user_content: str) -> str:
    template = getattr(tokenizer, "chat_template", None)
    if template:
        messages = [{"role": "user", "content": user_content}]
        return tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
        )
    return build_prompt_gemma_fallback(user_content)


def generate_text(
    model_id: str,
    user_content: str,
    max_new_tokens: int = 260,
) -> str:
    tokenizer, model = load_model_and_tokenizer(model_id)
    prompt = build_prompt(tokenizer, user_content)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    input_len = inputs["input_ids"].shape[1]

    output = model.generate(
        **inputs,
        max_new_tokens=max_new_tokens,
        temperature=0.2,
        do_sample=True,
        top_p=0.95,
        pad_token_id=tokenizer.pad_token_id,
        eos_token_id=tokenizer.eos_token_id,
    )
    new_tokens = output[0, input_len:]
    return tokenizer.decode(new_tokens, skip_special_tokens=True).strip()


def generate_abstract(
    model_id: str,
    title: str,
    introduction: str,
    conclusion: str,
    max_new_tokens: int = 260,
) -> str:
    user_content = build_user_content(title, introduction, conclusion)
    return generate_text(model_id=model_id, user_content=user_content, max_new_tokens=max_new_tokens)


def generate_summary_from_article(
    model_id: str,
    article_text: str,
    max_new_tokens: int = 260,
) -> str:
    user_content = build_user_content_from_article(article_text)
    return generate_text(model_id=model_id, user_content=user_content, max_new_tokens=max_new_tokens)