#实战项目二：自动摘要生成器

# 项目目标梳理
objectives = {
    "技术": ["TextRank/BERT抽取", "T5/BART/GPT生成", "ROUGE评估", "FastAPI服务"],
    "性能": ["ROUGE-1>0.4", "单条<2秒", "摘要10-20%原文长度"],
    "业务": ["多风格（简洁/技术）", "多长度文本（100-5000字）", "批量处理"]
}

# 精简版TextRank实现（去掉冗余，保留核心逻辑）
import numpy as np
import jieba
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.feature_extraction.text import TfidfVectorizer

class TextRank:
    def __init__(self, top_k=3):
        self.top_k = top_k
        self.damping = 0.85
        self.max_iter = 100
        self.tol = 1e-5

    def _split_sent(self, text):
        sents = []
        for s in text.replace('!', '。').replace('?', '。').split('。'):
            s = s.strip()
            if len(s) > 10: sents.append(s)
        return sents

    def _build_sim_mat(self, sents):
        # 中文TF-IDF（jieba分词）
        tfidf = TfidfVectorizer(tokenizer=jieba.lcut, token_pattern=None)
        tfidf_mat = tfidf.fit_transform(sents)
        sim_mat = cosine_similarity(tfidf_mat)
        np.fill_diagonal(sim_mat, 0)
        # 归一化
        row_sum = sim_mat.sum(axis=1, keepdims=True)
        row_sum[row_sum == 0] = 1
        return sim_mat / row_sum

    def summarize(self, text):
        sents = self._split_sent(text)
        if len(sents) <= self.top_k: return '。'.join(sents) + '。'
        
        sim_mat = self._build_sim_mat(sents)
        scores = np.ones(len(sents)) / len(sents)
        
        # PageRank迭代
        for _ in range(self.max_iter):
            new_scores = (1 - self.damping)/len(sents) + self.damping * sim_mat.T.dot(scores)
            if abs(new_scores - scores).sum() < self.tol: break
            scores = new_scores
        
        # 按原文顺序选Top-K
        top_idx = np.argsort(scores)[-self.top_k:][::-1]
        top_idx.sort()
        return '。'.join([sents[i] for i in top_idx]) + '。'

# 测试
text = "自然语言处理是AI的重要分支，研究人机语言交互。包含语音识别、文本分类、机器翻译等任务。近年来大模型取得突破，BERT和GPT是代表，ChatGPT在对话中表现出色。"
print(TextRank().summarize(text))

# 安装：pip install py-rouge
from rouge import Rouge

def eval_summary(candidate, reference):
    """简化版ROUGE评估（返回核心F1值）"""
    rouge = Rouge()
    scores = rouge.get_scores(candidate, reference)[0]
    return {
        "ROUGE-1-F1": round(scores['rouge-1']['f'], 4),
        "ROUGE-2-F1": round(scores['rouge-2']['f'], 4),
        "ROUGE-L-F1": round(scores['rouge-l']['f'], 4)
    }

# 测试
print(eval_summary("NLP是AI重要分支，大模型近年突破", "自然语言处理是人工智能重要分支，近年大模型取得突破性进展"))

# 安装：pip install transformers torch
from transformers import T5Tokenizer, T5ForConditionalGeneration
import torch

class T5Summary:
    def __init__(self):
        self.model_name = "t5-small"  # 入门用t5-small，效果可选t5-base/中文mT5
        self.tokenizer = T5Tokenizer.from_pretrained(self.model_name)
        self.model = T5ForConditionalGeneration.from_pretrained(self.model_name)
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model.to(self.device)

    def summarize(self, text, max_len=100, min_len=30):
        # T5要求前缀：summarize:
        input_text = "summarize: " + text
        inputs = self.tokenizer(input_text, return_tensors="pt", truncation=True, max_length=512).to(self.device)
        
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_length=max_len,
                min_length=min_len,
                num_beams=4,  # 束搜索，比随机采样更稳定
                no_repeat_ngram_size=2  # 防止重复
            )
        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)

# 测试
print(T5Summary().summarize(text))

class HybridSummary:
    def __init__(self):
        self.extract = TextRank(top_k=5)
        self.generate = T5Summary()

    def summarize(self, text, short_threshold=500):
        if len(text) < short_threshold:
            return self.extract.summarize(text)
        else:
            extracted = self.extract.summarize(text)
            return self.generate.summarize(extracted)

# app.py
from fastapi import FastAPI
from pydantic import BaseModel

app = FastAPI(title="自动摘要API", version="1.0")
hybrid = HybridSummary()

class SummaryReq(BaseModel):
    text: str
    short_threshold: int = 500

@app.post("/summary")
async def get_summary(req: SummaryReq):
    return {"summary": hybrid.summarize(req.text, req.short_threshold)}

# 启动命令：uvicorn app:app --host 0.0.0.0 --port 8000 --reload