Update api.py

api.py

@@ -1,54 +1,57 @@
-from flask import Blueprint, request, jsonify
-from transformers import DebertaV2Tokenizer, DebertaV2ForSequenceClassification
-import torch
-from utils import mask_pii
-
-api_bp = Blueprint("api", __name__)
-
-# Repo of Hugging Face Model Hub where Model is Pushed 
-REPO_ID = "Nikpatil/Email_classifier"
-MAX_LENGTH = 256
-
-tokenizer = DebertaV2Tokenizer.from_pretrained(REPO_ID)
-model = DebertaV2ForSequenceClassification.from_pretrained(REPO_ID)
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
-model.eval()
-
-id2label = {0: "Incident", 1: "Request", 2: "Problem", 3: "Change"}
-
-@api_bp.route("/classify", methods=["POST"])
-def classify_email():
-    data = request.get_json()
-    email_body = data.get("email_body", "")
-
-    if not email_body:
-        return jsonify({"Error": "Email body field is required"}), 400
-
-    masked_email, entities = mask_pii(email_body)
-
-    inputs = tokenizer(
-        masked_email,
-        add_special_tokens=True,
-        max_length=MAX_LENGTH,
-        padding='max_length',
-        truncation=True,
-        return_tensors='pt'
-    )
-    inputs = {k: v.to(device) for k, v in inputs.items()}
-
-    with torch.no_grad():
-        outputs = model(**inputs)
-
-    probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
-    predicted_class_id = torch.argmax(probs).item()
-    predicted_class = id2label[predicted_class_id]
-
-    return jsonify({
-        "input_email_body": email_body,
-        "list_of_masked_entities": entities,
-        "masked_email": masked_email,
-        "category_of_the_email": predicted_class
-    }), 200
-
-
+import os
+os.environ["HF_HOME"] = "/tmp/huggingface"
+
+from flask import Blueprint, request, jsonify
+from transformers import DebertaV2Tokenizer, DebertaV2ForSequenceClassification
+import torch
+from utils import mask_pii
+
+api_bp = Blueprint("api", __name__)
+
+# Repo of Hugging Face Model Hub where Model is Pushed 
+REPO_ID = "Nikpatil/Email_classifier"
+MAX_LENGTH = 256
+
+tokenizer = DebertaV2Tokenizer.from_pretrained(REPO_ID)
+model = DebertaV2ForSequenceClassification.from_pretrained(REPO_ID)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+model.eval()
+
+id2label = {0: "Incident", 1: "Request", 2: "Problem", 3: "Change"}
+
+@api_bp.route("/classify", methods=["POST"])
+def classify_email():
+    data = request.get_json()
+    email_body = data.get("email_body", "")
+
+    if not email_body:
+        return jsonify({"Error": "Email body field is required"}), 400
+
+    masked_email, entities = mask_pii(email_body)
+
+    inputs = tokenizer(
+        masked_email,
+        add_special_tokens=True,
+        max_length=MAX_LENGTH,
+        padding='max_length',
+        truncation=True,
+        return_tensors='pt'
+    )
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
+    predicted_class_id = torch.argmax(probs).item()
+    predicted_class = id2label[predicted_class_id]
+
+    return jsonify({
+        "input_email_body": email_body,
+        "list_of_masked_entities": entities,
+        "masked_email": masked_email,
+        "category_of_the_email": predicted_class
+    }), 200
+
+