sagar008 commited on
Commit
3d99217
·
verified ·
1 Parent(s): ff962a5

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +37 -0
app.py ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI
2
+ from pydantic import BaseModel
3
+ from transformers import pipeline, AutoTokenizer
4
+ import uvicorn
5
+
6
+ app = FastAPI()
7
+
8
+ MODEL_NAME = "VincentMuriuki/legal-summarizer"
9
+ summarizer = pipeline("summarization", model=MODEL_NAME)
10
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
11
+
12
+ class SummarizeInput(BaseModel):
13
+ text: str
14
+
15
+ class ChunkInput(BaseModel):
16
+ text: str
17
+ max_tokens: int = 512
18
+
19
+ @app.post("/summarize")
20
+ def summarize_text(data: SummarizeInput):
21
+ summary = summarizer(data.text, max_length=150, min_length=30, do_sample=False)
22
+ return {"summary": summary[0]["summary_text"]}
23
+
24
+ @app.post("/chunk")
25
+ def chunk_text(data: ChunkInput):
26
+ tokens = tokenizer.encode(data.text, truncation=False)
27
+ chunks = []
28
+
29
+ for i in range(0, len(tokens), data.max_tokens):
30
+ chunk_tokens = tokens[i:i + data.max_tokens]
31
+ chunk_text = tokenizer.decode(chunk_tokens, skip_special_tokens=True)
32
+ chunks.append(chunk_text.strip())
33
+
34
+ return {"chunks": chunks}
35
+
36
+ if _name_ == "_main_":
37
+ uvicorn.run(app, host="0.0.0.0", port=7860)