literature_summary.py

from transformers import PegasusTokenizer, PegasusForConditionalGeneration
import torch
import requests

USE_GEMINI_API = False

# === Pegasus Setup ===
if not USE_GEMINI_API:
    model_name = "google/pegasus-xsum"
    tokenizer = PegasusTokenizer.from_pretrained(model_name)
    model = PegasusForConditionalGeneration.from_pretrained(
        model_name, use_safetensors=True
    )
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device)
    MAX_INPUT_LENGTH = 512


def is_abstract_valid(abstract: str) -> bool:
    abstract_lower = abstract.lower()
    if len(abstract.split()) < 30:
        return False
    nonsense_phrases = [
        "apply for",
        "position",
        "job",
        "vacancy",
        "please apply",
        "contact",
        "email",
        "phone",
        "salary",
    ]
    for phrase in nonsense_phrases:
        if phrase in abstract_lower:
            return False
    return True


# === Pegasus Summarizer ===
def summarize_abstract_pegasus(abstract: str, max_length=60, num_beams=5):
    prompt = f"Summarize the following scientific abstract in 2-3 concise sentences:\n\n{abstract}\n\nSummary:"
    inputs = tokenizer(
        prompt, truncation=True, max_length=MAX_INPUT_LENGTH, return_tensors="pt"
    )
    input_ids = inputs.input_ids.to(device)
    attention_mask = inputs.attention_mask.to(device)

    with torch.no_grad():
        summary_ids = model.generate(
            input_ids=input_ids,
            attention_mask=attention_mask,
            max_length=max_length,
            num_beams=num_beams,
            early_stopping=True,
            no_repeat_ngram_size=3,
            repetition_penalty=2.0,
        )
    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
    return summary


def summarize_abstract_gemini(abstract: str):
    api_key = "AIzaSyA_hizLEmNNp7UXe3ID_fjLWj9bKsIK_lM"
    endpoint = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent?key={api_key}"
    headers = {"Content-Type": "application/json"}
    data = {
        "contents": [
            {
                "parts": [
                    {
                        "text": f"Summarize the following scientific abstract in 2-3 concise sentences:\n\n{abstract}"
                    }
                ]
            }
        ]
    }

    response = requests.post(endpoint, json=data, headers=headers)
    if response.status_code == 200:
        result = response.json()
        try:
            content = result["candidates"][0]["content"]
            # Extract the actual summary text
            summary_text = content["parts"][0]["text"]
            return summary_text
        except (KeyError, IndexError) as e:
            print(f"Unexpected API response structure: {e}")
            return "No summary returned (unexpected API response format)"
    else:
        print(f"Gemini API error: {response.status_code} - {response.text}")
        return "Error generating summary"


# === Wrapper to select method dynamically ===
def summarize_abstract(abstract: str):
    if USE_GEMINI_API:
        return summarize_abstract_gemini(abstract)
    else:
        return summarize_abstract_pegasus(abstract)


def summarize_literature_with_ref(papers):
    summaries = []
    reference_mapping = {}
    for idx, paper in enumerate(papers, start=1):
        abstract = paper.get("abstract", "").strip()
        title = paper.get("title", "Untitled").strip()
        if not abstract:
            continue
        if not is_abstract_valid(abstract):
            print(f"Skipping invalid abstract for paper: {title}")
            continue

        summary = summarize_abstract(abstract)
        ref_id = f"[{idx}]"
        reference_mapping[ref_id] = title
        summaries.append(f"{summary.strip()} {ref_id}")

    combined_summary = " ".join(summaries)  # single line with spaces

    references_text = "\n\nReferences:\n" + "\n".join(
        f"{ref_id} {title}" for ref_id, title in reference_mapping.items()
    )

    return combined_summary + "\n\n" + references_text