Update main.py
Browse files
main.py
CHANGED
@@ -13,13 +13,13 @@ from io import BytesIO
|
|
13 |
# -----------------------------------------------------------------------------
|
14 |
# CONFIGURATION
|
15 |
# -----------------------------------------------------------------------------
|
16 |
-
HUGGINGFACE_TOKEN = os.getenv("HF_TOKEN")
|
17 |
-
PORT
|
18 |
|
19 |
app = FastAPI(
|
20 |
-
title="AI
|
21 |
description="Backend for summarisation, captioning & QA",
|
22 |
-
version="1.2.
|
23 |
)
|
24 |
|
25 |
app.add_middleware(
|
@@ -31,7 +31,7 @@ app.add_middleware(
|
|
31 |
)
|
32 |
|
33 |
# -----------------------------------------------------------------------------
|
34 |
-
# OPTIONAL STATIC FILES
|
35 |
# -----------------------------------------------------------------------------
|
36 |
static_dir = Path("static")
|
37 |
if static_dir.exists():
|
@@ -40,8 +40,8 @@ if static_dir.exists():
|
|
40 |
# -----------------------------------------------------------------------------
|
41 |
# HUGGING FACE INFERENCE CLIENTS
|
42 |
# -----------------------------------------------------------------------------
|
43 |
-
summary_client
|
44 |
-
qa_client
|
45 |
image_caption_client = InferenceClient("nlpconnect/vit-gpt2-image-captioning", token=HUGGINGFACE_TOKEN)
|
46 |
|
47 |
# -----------------------------------------------------------------------------
|
@@ -57,7 +57,7 @@ def extract_text_from_docx(content: bytes) -> str:
|
|
57 |
return "\n".join(p.text for p in doc.paragraphs).strip()
|
58 |
|
59 |
def process_uploaded_file(file: UploadFile) -> str:
|
60 |
-
content
|
61 |
ext = file.filename.split(".")[-1].lower()
|
62 |
if ext == "pdf":
|
63 |
return extract_text_from_pdf(content)
|
@@ -73,36 +73,31 @@ def process_uploaded_file(file: UploadFile) -> str:
|
|
73 |
|
74 |
@app.get("/", response_class=HTMLResponse)
|
75 |
async def serve_index():
|
76 |
-
"""Return the frontend HTML page."""
|
77 |
return FileResponse("index.html")
|
78 |
|
79 |
# -------------------- Summarisation ------------------------------------------
|
80 |
-
|
81 |
@app.post("/api/summarize")
|
82 |
async def summarize_document(file: UploadFile = File(...)):
|
83 |
try:
|
84 |
text = process_uploaded_file(file)
|
85 |
if len(text) < 20:
|
86 |
return {"result": "Document too short to summarise."}
|
87 |
-
|
88 |
summary_raw = summary_client.summarization(text[:3000])
|
89 |
-
|
90 |
-
|
91 |
-
|
92 |
-
|
93 |
-
|
94 |
-
summary_txt = str(summary_raw)
|
95 |
-
|
96 |
return {"result": summary_txt}
|
97 |
except Exception as exc:
|
98 |
return JSONResponse(status_code=500, content={"error": f"Summarisation failure: {exc}"})
|
99 |
|
100 |
# -------------------- Image Caption -----------------------------------------
|
101 |
-
|
102 |
@app.post("/api/caption")
|
103 |
-
async def caption_image(
|
|
|
104 |
try:
|
105 |
-
img_bytes = await
|
106 |
img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
|
107 |
img.thumbnail((1024, 1024))
|
108 |
buf = BytesIO(); img.save(buf, format="JPEG")
|
@@ -118,29 +113,25 @@ async def caption_image(file: UploadFile = File(...)):
|
|
118 |
return JSONResponse(status_code=500, content={"error": f"Caption failure: {exc}"})
|
119 |
|
120 |
# -------------------- Question Answering ------------------------------------
|
121 |
-
|
122 |
@app.post("/api/qa")
|
123 |
async def question_answering(file: UploadFile = File(...), question: str = Form(...)):
|
124 |
try:
|
125 |
if file.content_type.startswith("image/"):
|
126 |
img_bytes = await file.read()
|
127 |
img = Image.open(io.BytesIO(img_bytes)).convert("RGB"); img.thumbnail((1024, 1024))
|
128 |
-
|
129 |
-
res = image_caption_client.image_to_text(
|
130 |
context = res.get("generated_text") if isinstance(res, dict) else str(res)
|
131 |
else:
|
132 |
context = process_uploaded_file(file)[:3000]
|
133 |
-
|
134 |
if not context:
|
135 |
return {"result": "No context β cannot answer."}
|
136 |
-
|
137 |
answer = qa_client.question_answering(question=question, context=context)
|
138 |
return {"result": answer.get("answer", "No answer found.")}
|
139 |
except Exception as exc:
|
140 |
return JSONResponse(status_code=500, content={"error": f"QA failure: {exc}"})
|
141 |
|
142 |
# -------------------- Health -------------------------------------------------
|
143 |
-
|
144 |
@app.get("/api/health")
|
145 |
async def health():
|
146 |
return {"status": "healthy", "hf_token_set": bool(HUGGINGFACE_TOKEN), "version": app.version}
|
@@ -148,7 +139,6 @@ async def health():
|
|
148 |
# -----------------------------------------------------------------------------
|
149 |
# ENTRYPOINT
|
150 |
# -----------------------------------------------------------------------------
|
151 |
-
|
152 |
if __name__ == "__main__":
|
153 |
import uvicorn
|
154 |
uvicorn.run(app, host="0.0.0.0", port=PORT)
|
|
|
13 |
# -----------------------------------------------------------------------------
|
14 |
# CONFIGURATION
|
15 |
# -----------------------------------------------------------------------------
|
16 |
+
HUGGINGFACE_TOKEN = os.getenv("HF_TOKEN")
|
17 |
+
PORT = int(os.getenv("PORT", 7860))
|
18 |
|
19 |
app = FastAPI(
|
20 |
+
title="AIβPowered WebβApp API",
|
21 |
description="Backend for summarisation, captioning & QA",
|
22 |
+
version="1.2.2",
|
23 |
)
|
24 |
|
25 |
app.add_middleware(
|
|
|
31 |
)
|
32 |
|
33 |
# -----------------------------------------------------------------------------
|
34 |
+
# OPTIONAL STATIC FILES
|
35 |
# -----------------------------------------------------------------------------
|
36 |
static_dir = Path("static")
|
37 |
if static_dir.exists():
|
|
|
40 |
# -----------------------------------------------------------------------------
|
41 |
# HUGGING FACE INFERENCE CLIENTS
|
42 |
# -----------------------------------------------------------------------------
|
43 |
+
summary_client = InferenceClient("facebook/bart-large-cnn", token=HUGGINGFACE_TOKEN)
|
44 |
+
qa_client = InferenceClient("deepset/roberta-base-squad2", token=HUGGINGFACE_TOKEN)
|
45 |
image_caption_client = InferenceClient("nlpconnect/vit-gpt2-image-captioning", token=HUGGINGFACE_TOKEN)
|
46 |
|
47 |
# -----------------------------------------------------------------------------
|
|
|
57 |
return "\n".join(p.text for p in doc.paragraphs).strip()
|
58 |
|
59 |
def process_uploaded_file(file: UploadFile) -> str:
|
60 |
+
content = file.file.read()
|
61 |
ext = file.filename.split(".")[-1].lower()
|
62 |
if ext == "pdf":
|
63 |
return extract_text_from_pdf(content)
|
|
|
73 |
|
74 |
@app.get("/", response_class=HTMLResponse)
|
75 |
async def serve_index():
|
|
|
76 |
return FileResponse("index.html")
|
77 |
|
78 |
# -------------------- Summarisation ------------------------------------------
|
|
|
79 |
@app.post("/api/summarize")
|
80 |
async def summarize_document(file: UploadFile = File(...)):
|
81 |
try:
|
82 |
text = process_uploaded_file(file)
|
83 |
if len(text) < 20:
|
84 |
return {"result": "Document too short to summarise."}
|
|
|
85 |
summary_raw = summary_client.summarization(text[:3000])
|
86 |
+
summary_txt = (
|
87 |
+
summary_raw[0].get("summary_text") if isinstance(summary_raw, list) else
|
88 |
+
summary_raw.get("summary_text") if isinstance(summary_raw, dict) else
|
89 |
+
str(summary_raw)
|
90 |
+
)
|
|
|
|
|
91 |
return {"result": summary_txt}
|
92 |
except Exception as exc:
|
93 |
return JSONResponse(status_code=500, content={"error": f"Summarisation failure: {exc}"})
|
94 |
|
95 |
# -------------------- Image Caption -----------------------------------------
|
|
|
96 |
@app.post("/api/caption")
|
97 |
+
async def caption_image(image: UploadFile = File(...)):
|
98 |
+
"""`image` field name matches frontend (was `file` before)."""
|
99 |
try:
|
100 |
+
img_bytes = await image.read()
|
101 |
img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
|
102 |
img.thumbnail((1024, 1024))
|
103 |
buf = BytesIO(); img.save(buf, format="JPEG")
|
|
|
113 |
return JSONResponse(status_code=500, content={"error": f"Caption failure: {exc}"})
|
114 |
|
115 |
# -------------------- Question Answering ------------------------------------
|
|
|
116 |
@app.post("/api/qa")
|
117 |
async def question_answering(file: UploadFile = File(...), question: str = Form(...)):
|
118 |
try:
|
119 |
if file.content_type.startswith("image/"):
|
120 |
img_bytes = await file.read()
|
121 |
img = Image.open(io.BytesIO(img_bytes)).convert("RGB"); img.thumbnail((1024, 1024))
|
122 |
+
buf = BytesIO(); img.save(buf, format="JPEG")
|
123 |
+
res = image_caption_client.image_to_text(buf.getvalue())
|
124 |
context = res.get("generated_text") if isinstance(res, dict) else str(res)
|
125 |
else:
|
126 |
context = process_uploaded_file(file)[:3000]
|
|
|
127 |
if not context:
|
128 |
return {"result": "No context β cannot answer."}
|
|
|
129 |
answer = qa_client.question_answering(question=question, context=context)
|
130 |
return {"result": answer.get("answer", "No answer found.")}
|
131 |
except Exception as exc:
|
132 |
return JSONResponse(status_code=500, content={"error": f"QA failure: {exc}"})
|
133 |
|
134 |
# -------------------- Health -------------------------------------------------
|
|
|
135 |
@app.get("/api/health")
|
136 |
async def health():
|
137 |
return {"status": "healthy", "hf_token_set": bool(HUGGINGFACE_TOKEN), "version": app.version}
|
|
|
139 |
# -----------------------------------------------------------------------------
|
140 |
# ENTRYPOINT
|
141 |
# -----------------------------------------------------------------------------
|
|
|
142 |
if __name__ == "__main__":
|
143 |
import uvicorn
|
144 |
uvicorn.run(app, host="0.0.0.0", port=PORT)
|