benkada commited on
Commit
8def51d
Β·
verified Β·
1 Parent(s): ef512e6

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +18 -28
main.py CHANGED
@@ -13,13 +13,13 @@ from io import BytesIO
13
  # -----------------------------------------------------------------------------
14
  # CONFIGURATION
15
  # -----------------------------------------------------------------------------
16
- HUGGINGFACE_TOKEN = os.getenv("HF_TOKEN") # set in HF Space secrets or env
17
- PORT = int(os.getenv("PORT", 7860)) # Spaces auto-set PORT; default 7860 locally
18
 
19
  app = FastAPI(
20
- title="AI-Powered Web-App API",
21
  description="Backend for summarisation, captioning & QA",
22
- version="1.2.1",
23
  )
24
 
25
  app.add_middleware(
@@ -31,7 +31,7 @@ app.add_middleware(
31
  )
32
 
33
  # -----------------------------------------------------------------------------
34
- # OPTIONAL STATIC FILES (only if ./static exists)
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 = 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,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 = file.file.read()
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
- if isinstance(summary_raw, list):
90
- summary_txt = summary_raw[0].get("summary_text", str(summary_raw))
91
- elif isinstance(summary_raw, dict):
92
- summary_txt = summary_raw.get("summary_text", str(summary_raw))
93
- else:
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(file: UploadFile = File(...)):
 
104
  try:
105
- img_bytes = await file.read()
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
- b = BytesIO(); img.save(b, format="JPEG")
129
- res = image_caption_client.image_to_text(b.getvalue())
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)