BASANT1896/text-summarisation

1from fastapi import FastAPI, Request, UploadFile, File
2from fastapi.responses import HTMLResponse
3from fastapi.templating import Jinja2Templates
4from pydantic import BaseModel
5from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
6import torch
7import re
8import gc  # Import garbage collector for manual cleanup
9from docx import Document
10from PyPDF2 import PdfReader
11import uvicorn
12import os
13
14# -----------------------------
15# FastAPI app
16# -----------------------------
17app = FastAPI(
18    title="Text Summarization System",
19    description="Summarize dialogues with T5!",
20    version="1.0"
21)
22
23# -----------------------------
24# Lazy Model Loading
25# -----------------------------
26model = None
27tokenizer = None
28# Force CPU device configuration entirely to save framework footprint overheads
29device = torch.device("cpu")
30
31
32def get_model():
33    global model, tokenizer
34    if model is None:
35        print("Loading model from Hugging Face Hub...")
36        model_id = "BASANT1896/saved_summary_model"
37
38        tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)
39
40        # OPTIMIZATION: Load tensors using 16-bit brain floating points (bfloat16) to cut footprint in half
41        model = AutoModelForSeq2SeqLM.from_pretrained(
42            model_id,
43            low_cpu_mem_usage=True,
44            torch_dtype=torch.bfloat16
45        )
46
47        model.to("cpu")
48        model.eval()
49    return model, tokenizer
50
51
52# -----------------------------
53# Templates
54# -----------------------------
55templates = Jinja2Templates(directory="templates")
56
57
58# -----------------------------
59# Input schema
60# -----------------------------
61class DialogueInput(BaseModel):
62    dialogue: str
63
64
65# -----------------------------
66# Utilities
67# -----------------------------
68# OPTIMIZATION: Reduce text cap from 6000 down to 2500 characters to prevent huge processing memory peaks
69MAX_INPUT_CHARS = 2500
70
71
72def clean_text(text: str) -> str:
73    text = re.sub(r'\r\n', ' ', text)
74    text = re.sub(r'\s+', ' ', text)
75    text = re.sub(r'<.*?>', '', text)
76    return text.strip()
77
78
79# -----------------------------
80# Summarization logic
81# -----------------------------
82def summarize_dialogue(dialogue: str) -> str:
83    # Free stray memory slots right before generation overhead starts
84    gc.collect()
85
86    current_model, current_tokenizer = get_model()  # Trigger lazy load
87
88    dialogue = clean_text(dialogue)
89    dialogue = dialogue[:MAX_INPUT_CHARS]
90    dialogue = "summarize: " + dialogue
91
92    inputs = current_tokenizer(
93        dialogue,
94        return_tensors="pt",
95        truncation=True,
96        max_length=512  # OPTIMIZATION: Reduced internal tokenization context max from 1024 to 512
97    )
98
99    input_ids = inputs["input_ids"].to(device)
100    attention_mask = inputs["attention_mask"].to(device)
101
102    input_length = input_ids.shape[1]
103
104    # OPTIMIZATION: Cap summary ranges tightly to restrict token tracking matrix memory usage
105    max_summary_len = max(60, int(input_length * 0.3))
106    max_summary_len = min(max_summary_len, 140)
107    min_summary_len = max(30, int(max_summary_len * 0.5))
108
109    with torch.no_grad():
110        outputs = current_model.generate(
111            input_ids,
112            attention_mask=attention_mask,
113            max_length=max_summary_len,
114            min_length=min_summary_len,
115            num_beams=2,  # OPTIMIZATION: Dropped search width from 4 down to 2 to minimize memory tracking tracks
116            length_penalty=1.1,
117            repetition_penalty=1.3,
118            no_repeat_ngram_size=3,
119            early_stopping=True
120        )
121
122    summary_text = current_tokenizer.decode(outputs[0], skip_special_tokens=True)
123
124    # Complete aggressive sweep instantly after text generation finishes
125    gc.collect()
126
127    return summary_text
128
129
130# -----------------------------
131# API Routes
132# -----------------------------
133@app.post("/summarize/")
134async def summarize(dialogue_input: DialogueInput):
135    summary = summarize_dialogue(dialogue_input.dialogue)
136    return {"summary": summary}
137
138
139@app.post("/upload")
140async def upload_file(file: UploadFile = File(...)):
141    text = ""
142    if file.filename.lower().endswith(".pdf"):
143        reader = PdfReader(file.file)
144        for page in reader.pages:
145            extracted = page.extract_text()
146            if extracted: text += extracted + " "
147    elif file.filename.lower().endswith(".docx"):
148        doc = Document(file.file)
149        for para in doc.paragraphs: text += para.text + " "
150    elif file.filename.lower().endswith(".txt"):
151        text = (await file.read()).decode("utf-8", errors="ignore")
152
153    text = clean_text(text)[:MAX_INPUT_CHARS]
154    return {"text": text}
155
156
157@app.get("/", response_class=HTMLResponse)
158async def home(request: Request):
159    # Fixed compatibility route signature structure for Python 3.14 + FastAPI
160    return templates.TemplateResponse(request, "index.html")
161
162
163if __name__ == "__main__":
164    port = int(os.environ.get("PORT", 8000))
165    uvicorn.run("app:app", host="0.0.0.0", port=port)