corderBackend / main.py
pius-code's picture
enhance summarize endpoint to dynamically adjust summary length based on input word count and return parameters used
2c871b7
raw
history blame
1.46 kB
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline
app = FastAPI()
summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
class TextInput(BaseModel):
text: str
@app.get("/")
async def root():
return {"message": "Welcome to the Text Summarization API!"}
@app.post("/summarize")
async def summarize_text(input: TextInput):
# Count approximate number of words (could be improved with tokenizer count in the future)
word_count = len(input.text.split())
if word_count < 50:
max_length = max(10, word_count // 2)
min_length = max(3, word_count // 4)
elif word_count < 200:
max_length = max(50, word_count // 3)
min_length = max(15, word_count // 6)
else:
max_length = max(100, word_count // 4)
min_length = max(30, word_count // 8)
# Prevent max_length from being too large (BART has token limits)
max_length = min(max_length, 1024)
# Generate summary with dynamic parameters
summary = summarizer(
input.text,
max_length=max_length,
min_length=min_length,
do_sample=True,
temperature=0.7,
num_beams=4
)
return {
"summary": summary[0]["summary_text"],
"parameters_used": {
"input_word_count": word_count,
"max_length": max_length,
"min_length": min_length
}
}