Add files via upload

KoBART.py

@@ -0,0 +1,52 @@
+from fastapi import FastAPI 
+from fastapi.middleware.cors import CORSMiddleware 
+from pydantic import BaseModel
+import torch 
+from transformers import PreTrainedTokenizerFast, BartForConditionalGeneration 
+from fastapi.responses import FileResponse
+
+tokenizer = PreTrainedTokenizerFast.from_pretrained('digit82/kobart-summarization')
+model = BartForConditionalGeneration.from_pretrained('digit82/kobart-summarization').to('cuda') 
+
+app = FastAPI()
+
+app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
+
+class TextRequest(BaseModel):
+    text: str
+
+
+@app.get("/")
+def get_root():
+    return FileResponse("index.html")
+
+@app.post("/summarize")
+async def summarize(req: TextRequest):
+    try:
+        text = req.text.replace('\n', ' ')
+        if len(text) < 100:
+            return {"error": f"입력 내용이 너무 짧습니다. 100글자 이상 입력해주세요. (현재 {len(text)}글자)"}
+        input_ids = tokenizer.encode(text)
+
+        if len(input_ids) > 1024:
+            return {
+                "error": f"입력 내용이 너무 깁니다. 최대 1024 단어까지 지원됩니다. (현재 {len(input_ids)} 단어)"
+            }
+
+        ids = [tokenizer.bos_token_id] + input_ids + [tokenizer.eos_token_id]
+
+        out = model.generate(
+            torch.tensor([ids]).to(model.device),
+            num_beams=10, 
+            max_length=120, 
+            min_length=50, 
+            length_penalty=1.2, 
+            repetition_penalty=2.5,
+            no_repeat_ngram_size=3, 
+            eos_token_id=tokenizer.eos_token_id 
+        )
+
+        return {"summary": tokenizer.decode(out.squeeze().tolist(), skip_special_tokens=True)}
+
+    except Exception as e:
+        return {"error": str(e)}