Added error handling

download_model.ipynb

@@ -59,12 +59,24 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 3,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Downloading config.json: 100%|██████████| 1.13k/1.13k [00:00<00:00, 283kB/s]\n",
+      "d:\\FYP\\lilt-app-without-fd\\lilt-env\\lib\\site-packages\\huggingface_hub\\file_download.py:133: UserWarning: `huggingface_hub` cache-system uses symlinks by default to efficiently store duplicated files but your machine does not support them in C:\\Users\\Gihantha Kavishka\\.cache\\huggingface\\hub. Caching files will still work but in a degraded version that might require more space on your disk. This warning can be disabled by setting the `HF_HUB_DISABLE_SYMLINKS_WARNING` environment variable. For more details, see https://huggingface.co/docs/huggingface_hub/how-to-cache#limitations.\n",
+      "To support symlinks on Windows, you either need to activate Developer Mode or to run Python as an administrator. In order to see activate developer mode, see this article: https://docs.microsoft.com/en-us/windows/apps/get-started/enable-your-device-for-development\n",
+      "  warnings.warn(message)\n",
+      "Downloading model.safetensors: 100%|██████████| 1.13G/1.13G [08:02<00:00, 2.35MB/s]\n"
+     ]
+    }
+   ],
    "source": [
     "# download the model\n",
-    "MODEL = \"pierreguillou/lilt-xlm-roberta-base-finetuned-funsd-iob-original\"\n",
+    "MODEL = \"kavg/LiLT-SER-Sin\"\n",
     "model = LiltForTokenClassification.from_pretrained(MODEL)\n",
     "\n",
     "# save the model\n",
@@ -83,28 +95,13 @@
    "cell_type": "code",
    "execution_count": 5,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Downloading config.json: 100%|██████████| 794/794 [00:00<00:00, 61.2kB/s]\n",
-      "d:\\FYP\\lilt-app-without-fd\\lilt-env\\lib\\site-packages\\huggingface_hub\\file_download.py:133: UserWarning: `huggingface_hub` cache-system uses symlinks by default to efficiently store duplicated files but your machine does not support them in C:\\Users\\Gihantha Kavishka\\.cache\\huggingface\\hub. Caching files will still work but in a degraded version that might require more space on your disk. This warning can be disabled by setting the `HF_HUB_DISABLE_SYMLINKS_WARNING` environment variable. For more details, see https://huggingface.co/docs/huggingface_hub/how-to-cache#limitations.\n",
-      "To support symlinks on Windows, you either need to activate Developer Mode or to run Python as an administrator. In order to see activate developer mode, see this article: https://docs.microsoft.com/en-us/windows/apps/get-started/enable-your-device-for-development\n",
-      "  warnings.warn(message)\n",
-      "Downloading pytorch_model.bin: 100%|██████████| 1.15G/1.15G [08:10<00:00, 2.34MB/s]\n",
-      "Some weights of the model checkpoint at kavg/layoutxlm-finetuned-xfund-fr-re were not used when initializing LiltModel: ['extractor.rel_classifier.linear.weight', 'extractor.entity_emb.weight', 'extractor.ffnn_tail.0.weight', 'extractor.ffnn_tail.3.bias', 'extractor.ffnn_head.3.weight', 'extractor.ffnn_head.0.weight', 'extractor.ffnn_tail.0.bias', 'extractor.ffnn_head.3.bias', 'extractor.rel_classifier.bilinear.weight', 'extractor.rel_classifier.linear.bias', 'extractor.ffnn_head.0.bias', 'extractor.ffnn_tail.3.weight']\n",
-      "- This IS expected if you are initializing LiltModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
-      "- This IS NOT expected if you are initializing LiltModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
-      "Some weights of LiltModel were not initialized from the model checkpoint at kavg/layoutxlm-finetuned-xfund-fr-re and are newly initialized: ['lilt.pooler.dense.bias', 'lilt.pooler.dense.weight']\n",
-      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
+    "from models import LiLTRobertaLikeForRelationExtraction\n",
+    "\n",
     "# download the model\n",
-    "MODEL =  'kavg/layoutxlm-finetuned-xfund-fr-re'\n",
-    "model = LiltModel.from_pretrained(MODEL)\n",
+    "MODEL =  'kavg/LiLT-RE-IT-Sin'\n",
+    "model = LiLTRobertaLikeForRelationExtraction.from_pretrained(MODEL)\n",
     "\n",
     "# save the model\n",
     "save_dir = \"models/lilt-re\"\n",

main.py

@@ -5,7 +5,7 @@ from PIL import Image
 from transformers import LiltForTokenClassification, AutoTokenizer
 import token_classification
 import torch
-from fastapi import FastAPI, UploadFile, Form
+from fastapi import FastAPI, UploadFile, Form, HTTPException
 from contextlib import asynccontextmanager
 import json
 import io
@@ -32,25 +32,49 @@ app = FastAPI(lifespan=lifespan)
 @app.post("/submit-doc")
 async def ProcessDocument(file: UploadFile):
   content = await file.read()
-  tokenClassificationOutput, ocr_df, img_size = LabelTokens(content)
-  reOutput = ExtractRelations(tokenClassificationOutput, ocr_df, img_size)
+  ocr_df, image = ApplyOCR(content)
+  if len(ocr_df) < 2:
+    raise HTTPException(status_code=400, detail="Cannot apply OCR to the image")
+  try:
+    tokenClassificationOutput, img_size = LabelTokens(ocr_df, image)
+    reOutput = ExtractRelations(tokenClassificationOutput, ocr_df, img_size)
+  except:
+    raise HTTPException(status_code=400, detail="Invalid Image")
   return reOutput
 
 @app.post("/submit-doc-base64")
 async def ProcessDocument(file: str = Form(...)):
-  head, file = file.split(',')
-  str_as_bytes = str.encode(file)
-  content = b64decode(str_as_bytes)
-  tokenClassificationOutput, ocr_df, img_size = LabelTokens(content)
-  reOutput = ExtractRelations(tokenClassificationOutput, ocr_df, img_size)
+  try:
+    head, file = file.split(',')
+    str_as_bytes = str.encode(file)
+    content = b64decode(str_as_bytes)
+  except:
+    raise HTTPException(status_code=400, detail="Invalid image")
+  ocr_df, image = ApplyOCR(content)
+  if len(ocr_df) < 2:
+    raise HTTPException(status_code=400, detail="Cannot apply OCR to the image")
+  try:
+    tokenClassificationOutput, img_size = LabelTokens(ocr_df, image)
+    reOutput = ExtractRelations(tokenClassificationOutput, ocr_df, img_size)
+  except:
+    raise HTTPException(status_code=400, detail="Invalid Image")
   return reOutput
 
-def LabelTokens(content):
-  image = Image.open(io.BytesIO(content))
-  ocr_df = config['vision_client'].ocr(content, image)
+def ApplyOCR(content):
+  try:
+    image = Image.open(io.BytesIO(content))
+  except:
+    raise HTTPException(status_code=400, detail="Invalid image")
+  try:
+    ocr_df = config['vision_client'].ocr(content, image)
+  except:
+    raise HTTPException(status_code=400, detail="OCR process failed")
+  return ocr_df, image
+
+def LabelTokens(ocr_df, image):
   input_ids, attention_mask, token_type_ids, bbox, token_actual_boxes, offset_mapping = config['processor'].process(ocr_df, image = image)
   token_labels = token_classification.classifyTokens(config['ser_model'], input_ids, attention_mask, bbox, offset_mapping)
-  return {"token_labels": token_labels, "input_ids": input_ids, "bbox":bbox, "attention_mask":attention_mask}, ocr_df, image.size
+  return {"token_labels": token_labels, "input_ids": input_ids, "bbox":bbox, "attention_mask":attention_mask}, image.size
 
 def ExtractRelations(tokenClassificationOutput, ocr_df, img_size):
   token_labels = tokenClassificationOutput['token_labels']