Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
@@ -1,25 +1,25 @@
|
|
1 |
from fastapi import FastAPI
|
2 |
from pydantic import BaseModel
|
3 |
from transformers import pipeline, AutoTokenizer
|
4 |
-
import uvicorn
|
5 |
import os
|
|
|
6 |
app = FastAPI()
|
7 |
-
os.environ['TRANSFORMERS_CACHE'] = '/app/.cache'
|
8 |
-
os.environ['HF_HOME'] = '/app/.cache'
|
9 |
-
os.makedirs('/app/.cache', exist_ok=True)
|
10 |
-
HF_AUTH_TOKEN=os.environ.get("HF_TOKEN")
|
11 |
|
|
|
|
|
|
|
|
|
12 |
|
13 |
MODEL_NAME = "VincentMuriuki/legal-summarizer"
|
|
|
14 |
summarizer = pipeline("summarization", model=MODEL_NAME, use_auth_token=HF_AUTH_TOKEN)
|
15 |
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=HF_AUTH_TOKEN)
|
16 |
-
|
17 |
class SummarizeInput(BaseModel):
|
18 |
text: str
|
19 |
|
20 |
class ChunkInput(BaseModel):
|
21 |
text: str
|
22 |
-
max_tokens: int = 512
|
23 |
|
24 |
@app.post("/summarize")
|
25 |
def summarize_text(data: SummarizeInput):
|
@@ -37,6 +37,3 @@ def chunk_text(data: ChunkInput):
|
|
37 |
chunks.append(chunk_text.strip())
|
38 |
|
39 |
return {"chunks": chunks}
|
40 |
-
|
41 |
-
if __name__ == "__main__":
|
42 |
-
uvicorn.run(app, host="0.0.0.0", port=7860)
|
|
|
1 |
from fastapi import FastAPI
|
2 |
from pydantic import BaseModel
|
3 |
from transformers import pipeline, AutoTokenizer
|
|
|
4 |
import os
|
5 |
+
|
6 |
app = FastAPI()
|
|
|
|
|
|
|
|
|
7 |
|
8 |
+
os.environ["TRANSFORMERS_CACHE"] = "/home/user/.cache/huggingface"
|
9 |
+
os.environ["HF_HOME"] = "/home/user/.cache/huggingface"
|
10 |
+
|
11 |
+
HF_AUTH_TOKEN = os.environ.get("HF_TOKEN")
|
12 |
|
13 |
MODEL_NAME = "VincentMuriuki/legal-summarizer"
|
14 |
+
|
15 |
summarizer = pipeline("summarization", model=MODEL_NAME, use_auth_token=HF_AUTH_TOKEN)
|
16 |
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=HF_AUTH_TOKEN)
|
|
|
17 |
class SummarizeInput(BaseModel):
|
18 |
text: str
|
19 |
|
20 |
class ChunkInput(BaseModel):
|
21 |
text: str
|
22 |
+
max_tokens: int = 512
|
23 |
|
24 |
@app.post("/summarize")
|
25 |
def summarize_text(data: SummarizeInput):
|
|
|
37 |
chunks.append(chunk_text.strip())
|
38 |
|
39 |
return {"chunks": chunks}
|
|
|
|
|
|