Spaces:
Running
Running
| from fastapi import FastAPI, HTTPException | |
| from pydantic import BaseModel | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from transformers import AutoTokenizer, AutoModelForSeq2SeqLM | |
| import os | |
| import gc | |
| import torch | |
| import traceback | |
| app = FastAPI() | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| token = os.environ.get("HF_TOKEN") | |
| class SummarizeRequest(BaseModel): | |
| text: str | |
| model_type: str | |
| def read_root(): | |
| return {"message": "Text Summarizer API Backend is Running!"} | |
| def summarize(req: SummarizeRequest): | |
| models = { | |
| "bart": "kUrrooYuki/bart-large-cnn-text-summarize", | |
| "t5": "kUrrooYuki/t5-base-text-summarize", | |
| "led": "kUrrooYuki/led-base-text-summarize" | |
| } | |
| if req.model_type not in models: | |
| raise HTTPException(status_code=400, detail="Model not recognized") | |
| try: | |
| model_repo = models[req.model_type] | |
| print(f"-> [1/3] Memuat {model_repo} ke RAM...", flush=True) | |
| tokenizer = AutoTokenizer.from_pretrained(model_repo, token=token) | |
| model = AutoModelForSeq2SeqLM.from_pretrained(model_repo, token=token) | |
| text_input = req.text | |
| if req.model_type == "t5": | |
| text_input = "summarize: " + text_input | |
| print("-> [2/3] Starting the AI computation process (please wait)...", flush=True) | |
| inputs = tokenizer(text_input, return_tensors="pt", max_length=1024, truncation=True) | |
| in_ids = inputs["input_ids"] | |
| att_mask = inputs["attention_mask"] | |
| if not torch.is_tensor(in_ids): | |
| if isinstance(in_ids, list) and (len(in_ids) == 0 or not isinstance(in_ids[0], list)): | |
| in_ids = [in_ids] | |
| in_ids = torch.tensor(in_ids) | |
| if not torch.is_tensor(att_mask): | |
| if isinstance(att_mask, list) and (len(att_mask) == 0 or not isinstance(att_mask[0], list)): | |
| att_mask = [att_mask] | |
| att_mask = torch.tensor(att_mask) | |
| beams_config = 2 if req.model_type == "bart" else 1 | |
| ngram_config = 3 if req.model_type == "bart" else 0 | |
| summary_ids = model.generate( | |
| input_ids=in_ids, | |
| attention_mask=att_mask, | |
| max_new_tokens=100, | |
| min_length=20, | |
| num_beams=beams_config, | |
| no_repeat_ngram_size=ngram_config, | |
| early_stopping=True | |
| ) | |
| print("-> [3/3] Computation complete! Translating tokens to text...", flush=True) | |
| summary_result = tokenizer.decode(summary_ids[0], skip_special_tokens=True) | |
| summary_result = summary_result.strip() | |
| if summary_result.lower().startswith("a's "): | |
| summary_result = summary_result[4:].strip() | |
| elif summary_result.lower().startswith("a "): | |
| summary_result = summary_result[2:].strip() | |
| if len(summary_result) > 0: | |
| summary_result = summary_result[0].upper() + summary_result[1:] | |
| del tokenizer | |
| del model | |
| gc.collect() | |
| print("-> Success. The RAM has been cleared.", flush=True) | |
| return {"status": "success", "summary": summary_result} | |
| except Exception as e: | |
| print("=== FULL ERROR TRACEBACK ===", flush=True) | |
| traceback.print_exc() | |
| print("============================", flush=True) | |
| raise HTTPException(status_code=500, detail=str(e)) |