sagar008 commited on
Commit
01728a1
·
verified ·
1 Parent(s): c47c86d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +7 -10
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}