Update app.py

app.py

@@ -1,19 +1,28 @@
+import os
 from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import pipeline, AutoTokenizer
-import os
+import uvicorn
 
-app = FastAPI()
+# Create a custom cache directory inside app folder
+os.environ["TRANSFORMERS_CACHE"] = "/app/hf_cache"
 
-os.environ["TRANSFORMERS_CACHE"] = "/data/cache"
-os.environ["HF_HOME"] = "/data/cache"
+app = FastAPI()
 
 HF_AUTH_TOKEN = os.environ.get("HF_TOKEN")
-
 MODEL_NAME = "VincentMuriuki/legal-summarizer"
 
-summarizer = pipeline("summarization", model=MODEL_NAME, use_auth_token=HF_AUTH_TOKEN)
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=HF_AUTH_TOKEN)
+summarizer = pipeline(
+    "summarization",
+    model=MODEL_NAME,
+    use_auth_token=HF_AUTH_TOKEN
+)
+
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_NAME,
+    use_auth_token=HF_AUTH_TOKEN
+)
+
 class SummarizeInput(BaseModel):
     text: str
 
@@ -37,3 +46,6 @@ def chunk_text(data: ChunkInput):
         chunks.append(chunk_text.strip())
 
     return {"chunks": chunks}
+
+if _name_ == "_main_":
+    uvicorn.run(app, host="0.0.0.0", port=7860)