init commit

Dockerfile

@@ -0,0 +1,13 @@
+FROM python:3.8
+
+WORKDIR /code
+
+COPY ./requirements.txt /code/requirements.txt
+
+RUN pip install --upgrade -r /code/requirements.txt
+
+COPY ./model_finetuned_clear.pkl ./model_finetuned_clear.pkl
+COPY ./required_classes.py ./required_classes.py
+COPY ./app.py ./app.py
+
+CMD ["python", "app.py"]

app.py

@@ -0,0 +1,67 @@
+print('INFO: import modules')
+from flask import Flask, request
+import json
+import pickle
+import numpy as np
+
+from required_classes import BertEmbedder, PredictModel
+
+
+print('INFO: loading model')
+try:
+    with open('model_finetuned_clear.pkl', 'rb') as f:
+        model = pickle.load(f)
+    model.batch_size = 1
+    print('INFO: model loaded')
+except Exception as e:
+    print(f"ERROR: loading models failed with: {str(e)}")
+
+def classify_code(text, top_n):
+    embed = model._texts2vecs([text])
+    probs = model.classifier_code.predict_proba(embed)
+    best_n = np.flip(np.argsort(probs, axis=1,)[0,-top_n:])
+    preds = [{'code': model.classifier_code.classes_[i], 'proba': probs[0][i]} for i in best_n]
+    return preds
+
+def classify_group(text, top_n):
+    embed = model._texts2vecs([text])
+    probs = model.classifier_group.predict_proba(embed)
+    best_n = np.flip(np.argsort(probs, axis=1,)[0,-top_n:])
+    preds = [{'group': model.classifier_group.classes_[i], 'proba': probs[0][i]} for i in best_n]
+    return preds
+
+
+app = Flask(__name__)
+
+@app.get("/")
+def test_get():
+    return {'hello': 'world'}
+
+@app.route("/test", methods=['POST'])
+def test():
+    data = request.form
+    return {'response': data}
+
+@app.route("/predict", methods=['POST'])
+def read_root(): 
+    data = request.form
+    text = str(data['text'])
+    top_n = int(data['top_n'])
+
+    if top_n < 1:
+        return {'error': 'top_n should be geather than 0'}
+    if text.strip() == '':
+        return {'error': 'text is empty'}
+    
+    pred_codes = classify_code(text, top_n)
+    pred_groups = classify_group(text, top_n)
+    result = {
+        "icd10": 
+            {'result': pred_codes[0]['code'], 'details': pred_codes},
+        "dx_group":
+            {'result': pred_groups[0]['group'], 'details': pred_groups}
+    }
+    return result
+
+if __name__ == "__main__":
+    app.run(host='0.0.0.0', port=7860)

model_finetuned_clear.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4c40076019c4b4767021bf208200a8104f0910669d0b56952e6b2eb62b1539d3
+size 434856921

required_classes.py

@@ -0,0 +1,74 @@
+import numpy as np
+from typing import List
+
+
+class BertEmbedder:
+    def __init__(self, model_path:str, cut_head:bool=False):
+        """
+            cut_head = True if the model have classifier head
+        """
+        self.embedder = BertForSequenceClassification.from_pretrained(model_path)
+        self.max_length = self.embedder.config.max_position_embeddings
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path, max_length=self.max_length)
+
+        if cut_head:
+            self.embedder = self.embedder.bert
+
+        self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        self.embedder.to(self.device)
+
+    def __call__(self, text: str):
+        encoded_input = self.tokenizer(text, 
+                                       return_tensors='pt', 
+                                       max_length=self.max_length,
+                                       padding=True,
+                                       truncation=True).to(self.device)
+        model_output = self.embedder(**encoded_input)
+        text_embed = model_output.pooler_output[0].cpu()
+        return text_embed
+
+    def batch_predict(self, texts: List[str]):
+        encoded_input = self.tokenizer(texts, 
+                                       return_tensors='pt', 
+                                       max_length=self.max_length,
+                                       padding=True,
+                                       truncation=True).to(self.device)
+        model_output = self.embedder(**encoded_input)
+        texts_embeds = model_output.pooler_output.cpu()
+        return texts_embeds
+
+class PredictModel:
+    def __init__(self, embedder, classifier, batch_size=8):
+        self.batch_size = batch_size
+        self.embedder = embedder
+        self.classifier = classifier
+
+    def _texts2vecs(self, texts, log=False):
+        embeds = []
+        batches_texts = np.array_split(texts, len(texts) // self.batch_size)
+        if log:
+            iterator = tqdm(batches_texts)
+        else:
+            iterator = batches_texts
+        for batch_texts in iterator:
+            batch_texts = batch_texts.tolist()
+            embeds += self.embedder.batch_predict(batch_texts).tolist()
+        embeds = np.array(embeds)
+        return embeds
+
+    def fit(self, texts: List[str], labels: List[str], log: bool=False):
+        if log:
+            print('Start text2vec transform')
+        embeds = self._texts2vecs(texts, log)
+        if log:
+            print('Start classifier fitting')
+        self.classifier.fit(embeds, labels)
+
+    def predict(self, texts: List[str], log: bool=False):
+        if log:
+            print('Start text2vec transform')
+        embeds = self._texts2vecs(texts, log)
+        if log:
+            print('Start classifier prediction')
+        prediction = self.classifier.predict(embeds)
+        return prediction

requirements.txt

@@ -0,0 +1,5 @@
+numpy==1.22.4
+torch==2.0.1
+scikit-learn==1.2.2
+transformers==4.29.2
+flask==2.0.3