Spaces:
Sleeping
Sleeping
File size: 1,019 Bytes
dc08424 d5da2a8 49332a9 bbe006e 49332a9 bbe006e dc08424 f4df259 bbe006e dc08424 f4df259 dc08424 30a81ee dc08424 bbe006e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 |
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch
import uvicorn
# Create FastAPI app
app = FastAPI()
# Load the tokenizer and model
MODEL_NAME = "facebook/bart-large-cnn" # A lightweight summarization model
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME).to("cpu") # Use "cuda" if you have a GPU
# Define input format
class InputText(BaseModel):
text: str
@app.post("/summarize")
async def summarize_text(input_text: InputText):
inputs = tokenizer(input_text.text, return_tensors="pt", max_length=1024, truncation=True)
summary_ids = model.generate(inputs.input_ids, max_length=150, min_length=50, length_penalty=2.0)
summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
return {"summary": summary}
# Ensure the application starts when running locally
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=7860)
|