AI Engine API

main.py

@@ -0,0 +1,39 @@
+from fastapi import FastAPI, UploadFile
+from fastapi.middleware.cors import CORSMiddleware
+
+import utils.handle as handle
+
+app = FastAPI()
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+@app.post("/training_data_from_utags_json")
+async def training_data_from_utags_json(file: UploadFile, savePath: str):
+    return handle.trainingDataFromUTagsJSON(file, savePath)
+
+
+@app.post("/training_data_from_prompts_for_bert")
+async def training_data_from_prompts_for_bert(file: UploadFile, savePath: str):
+    return handle.trainingDataFromPromptsForBERT(file, savePath)
+
+
+@app.post("/augment_data_using_vector_space_algorithm")
+async def augment_data_using_vector_space_algorithm(file: UploadFile, savePath: str):
+    return handle.augmentDataUsingVectorSpaceAlgorithm(file, savePath)
+
+
+@app.post("/get_symptoms_causes_and_disease_name_from_json")
+async def get_symptoms_causes_and_disease_name_from_json(file: UploadFile, savePath: str):
+    return handle.getSymptomsCausesAndDiseaseNameFromJSON(file, savePath)
+
+
+@app.post("/train_model_on_sagemaker")
+async def train_model_on_sagemaker(trainDataPath: str, testDataPath: str, file: UploadFile | None = None):
+    return handle.trainModelOnSageMaker(trainDataPath, testDataPath, file)

requirements.in

@@ -1 +1,10 @@
-gradio
+boto3
+datasets
+fastapi
+gradio
+sagemaker
+scikit-learn
+seaborn
+torch
+transformers
+uvicorn

requirements.txt

@@ -7,7 +7,10 @@
 aiofiles==23.1.0
     # via gradio
 aiohttp==3.8.4
-    # via gradio
+    # via
+    #   datasets
+    #   fsspec
+    #   gradio
 aiosignal==1.3.1
     # via aiohttp
 altair==4.2.2
@@ -18,10 +21,19 @@ anyio==3.6.2
     #   starlette
 async-timeout==4.0.2
     # via aiohttp
-attrs==23.1.0
+attrs==22.2.0
     # via
     #   aiohttp
     #   jsonschema
+    #   sagemaker
+boto3==1.26.133
+    # via
+    #   -r requirements.in
+    #   sagemaker
+botocore==1.29.133
+    # via
+    #   boto3
+    #   s3transfer
 certifi==2022.12.7
     # via
     #   httpcore
@@ -33,32 +45,51 @@ charset-normalizer==3.1.0
     #   requests
 click==8.1.3
     # via uvicorn
+cloudpickle==2.2.1
+    # via sagemaker
 colorama==0.4.6
     # via
     #   click
     #   tqdm
+contextlib2==21.6.0
+    # via schema
 contourpy==1.0.7
     # via matplotlib
 cycler==0.11.0
     # via matplotlib
+datasets==2.12.0
+    # via -r requirements.in
+dill==0.3.6
+    # via
+    #   datasets
+    #   multiprocess
+    #   pathos
 entrypoints==0.4
     # via altair
 fastapi==0.95.1
-    # via gradio
+    # via
+    #   -r requirements.in
+    #   gradio
 ffmpy==0.3.0
     # via gradio
 filelock==3.12.0
-    # via huggingface-hub
+    # via
+    #   huggingface-hub
+    #   torch
+    #   transformers
 fonttools==4.39.3
     # via matplotlib
 frozenlist==1.3.3
     # via
     #   aiohttp
     #   aiosignal
-fsspec==2023.4.0
+fsspec[http]==2023.4.0
     # via
+    #   datasets
     #   gradio-client
     #   huggingface-hub
+google-pasta==0.2.0
+    # via sagemaker
 gradio==3.28.3
     # via -r requirements.in
 gradio-client==0.2.0
@@ -75,20 +106,33 @@ httpx==0.24.0
     #   gradio-client
 huggingface-hub==0.14.1
     # via
+    #   datasets
     #   gradio
     #   gradio-client
+    #   transformers
 idna==3.4
     # via
     #   anyio
     #   httpx
     #   requests
     #   yarl
+importlib-metadata==4.13.0
+    # via sagemaker
 jinja2==3.1.2
     # via
     #   altair
     #   gradio
+    #   torch
+jmespath==1.0.1
+    # via
+    #   boto3
+    #   botocore
+joblib==1.2.0
+    # via scikit-learn
 jsonschema==4.17.3
-    # via altair
+    # via
+    #   altair
+    #   sagemaker
 kiwisolver==1.4.4
     # via matplotlib
 linkify-it-py==2.0.2
@@ -102,37 +146,76 @@ markupsafe==2.1.2
     #   gradio
     #   jinja2
 matplotlib==3.7.1
-    # via gradio
+    # via
+    #   gradio
+    #   seaborn
 mdit-py-plugins==0.3.3
     # via gradio
 mdurl==0.1.2
     # via markdown-it-py
+mpmath==1.3.0
+    # via sympy
 multidict==6.0.4
     # via
     #   aiohttp
     #   yarl
+multiprocess==0.70.14
+    # via
+    #   datasets
+    #   pathos
+networkx==3.1
+    # via torch
 numpy==1.24.3
     # via
     #   altair
     #   contourpy
+    #   datasets
     #   gradio
     #   matplotlib
     #   pandas
+    #   pyarrow
+    #   sagemaker
+    #   scikit-learn
+    #   scipy
+    #   seaborn
+    #   transformers
 orjson==3.8.11
     # via gradio
 packaging==23.1
     # via
+    #   datasets
     #   gradio-client
     #   huggingface-hub
     #   matplotlib
+    #   sagemaker
+    #   transformers
 pandas==2.0.1
     # via
     #   altair
+    #   datasets
     #   gradio
+    #   sagemaker
+    #   seaborn
+pathos==0.3.0
+    # via sagemaker
 pillow==9.5.0
     # via
     #   gradio
     #   matplotlib
+platformdirs==3.5.1
+    # via sagemaker
+pox==0.3.2
+    # via pathos
+ppft==1.7.6.6
+    # via pathos
+protobuf==3.20.3
+    # via
+    #   protobuf3-to-dict
+    #   sagemaker
+protobuf3-to-dict==0.1.5
+    # via sagemaker
+pyarrow==12.0.0
+    # via datasets
 pydantic==1.10.7
     # via
     #   fastapi
@@ -147,25 +230,54 @@ pyrsistent==0.19.3
     # via jsonschema
 python-dateutil==2.8.2
     # via
+    #   botocore
     #   matplotlib
     #   pandas
 python-multipart==0.0.6
     # via gradio
 pytz==2023.3
     # via pandas
-pyyaml==6.0
+pyyaml==5.4.1
     # via
+    #   datasets
     #   gradio
     #   huggingface-hub
+    #   sagemaker
+    #   transformers
+regex==2023.5.5
+    # via transformers
 requests==2.30.0
     # via
+    #   datasets
+    #   fsspec
     #   gradio
     #   gradio-client
     #   huggingface-hub
+    #   responses
+    #   transformers
+responses==0.18.0
+    # via datasets
+s3transfer==0.6.1
+    # via boto3
+sagemaker==2.154.0
+    # via -r requirements.in
+schema==0.7.5
+    # via sagemaker
+scikit-learn==1.2.2
+    # via -r requirements.in
+scipy==1.10.1
+    # via scikit-learn
+seaborn==0.12.2
+    # via -r requirements.in
 semantic-version==2.10.0
     # via gradio
 six==1.16.0
-    # via python-dateutil
+    # via
+    #   google-pasta
+    #   protobuf3-to-dict
+    #   python-dateutil
+smdebug-rulesconfig==1.0.1
+    # via sagemaker
 sniffio==1.3.0
     # via
     #   anyio
@@ -173,27 +285,52 @@ sniffio==1.3.0
     #   httpx
 starlette==0.26.1
     # via fastapi
+sympy==1.12
+    # via torch
+tblib==1.7.0
+    # via sagemaker
+threadpoolctl==3.1.0
+    # via scikit-learn
+tokenizers==0.13.3
+    # via transformers
 toolz==0.12.0
     # via altair
+torch==2.0.1
+    # via -r requirements.in
 tqdm==4.65.0
-    # via huggingface-hub
+    # via
+    #   datasets
+    #   huggingface-hub
+    #   transformers
+transformers==4.29.1
+    # via -r requirements.in
 typing-extensions==4.5.0
     # via
     #   gradio
     #   gradio-client
     #   huggingface-hub
     #   pydantic
+    #   torch
 tzdata==2023.3
     # via pandas
 uc-micro-py==1.0.2
     # via linkify-it-py
-urllib3==2.0.2
-    # via requests
+urllib3==1.26.15
+    # via
+    #   botocore
+    #   requests
+    #   responses
 uvicorn==0.22.0
-    # via gradio
+    # via
+    #   -r requirements.in
+    #   gradio
 websockets==11.0.2
     # via
     #   gradio
     #   gradio-client
+xxhash==3.2.0
+    # via datasets
 yarl==1.9.2
     # via aiohttp
+zipp==3.15.0
+    # via importlib-metadata

tasks/data/dataAugmentation.py

@@ -0,0 +1,66 @@
+import pandas as pd
+
+
+def augmentDataWithVectorSpaceAlgorithm(data: pd.DataFrame) -> pd.DataFrame:
+    """
+    Augment the Data
+    =================
+    Parameters:
+    -----------
+        data:
+            description: Data to augment
+            type: pd.DataFrame
+    -----------
+    Returns:
+    --------
+        data:
+            description: Augmented data
+            type: pd.DataFrame
+    --------------------------------------------------------------------------------------------
+    Working:
+    --------
+        - Create a DataFrame from data
+        - Remove a symptom or cause from the new DataFrame
+        - Check if the resulting row is present in the original data
+        - If not present, add the resulting row to the new DataFrame
+        - Repeat steps 1-3 for all symptoms and causes
+        - Remove the rows with sum = 0
+        - Remove the same rows from the new DataFrame
+        - Add the new DataFrame to the original data
+        - Return the Resulting DataFrame
+    --------------------------------------------------------------------------------------------
+    """
+    
+    # Get the number of columns with symptoms_ prefix
+    numberOfSymptoms = len([col for col in data.columns if col.startswith('symptoms_')])
+    
+    symptoms = data.columns[1:numberOfSymptoms]
+    causes = data.columns[numberOfSymptoms:]   
+    df = data
+    
+    for index, row in data.iterrows():
+        for symptom in symptoms:
+            if row[symptom] == 1: # type: ignore
+                row[symptom] = 0
+                df = df.append(row, ignore_index=True) # type: ignore
+                row[symptom] = 1
+                df.append(row, ignore_index=True)
+        
+        for cause in causes:
+            if row[cause] == 1:   # type: ignore
+                row[cause] = 0
+                df = df.append(row, ignore_index=True) # type: ignore
+                row[cause] = 1
+                df.append(row, ignore_index=True)
+    
+    
+    print(f"data before drop_duplicates: {df}")
+    
+    df = df[(df.sum(axis=1, numeric_only=True) != 0)]
+    data = data.append(df, ignore_index=True) # type: ignore
+    data = data.drop_duplicates(subset=df.columns.difference(['disease']), keep=False)
+    data.reset_index(drop=True, inplace=True)
+    
+    print(f"final data: {data}")
+    
+    return data

tasks/data/dataEngineering.py

@@ -0,0 +1,74 @@
+import pandas as pd
+from sklearn.preprocessing import MultiLabelBinarizer
+
+
+def trainingDataFromUTagsJSON(data: dict) -> pd.DataFrame:
+    """
+    Get the training data from the UTags JSON file
+    ==============================================
+    Parameters:
+    -----------
+        data:
+            description: UTags JSON file
+            type: dict
+    -----------
+    Returns:
+    --------
+        data:
+            description: Training data
+            type: pd.DataFrame
+    """
+    df = pd.DataFrame()
+    
+    df['disease'] = [disease.disease_persian[0] for disease in data['diseases']] # disease[UTag]
+    df['symptoms'] = [disease.symptom_eng for disease in data['diseases']]
+    df['causes'] = [disease.cause_eng for disease in data['diseases']]
+    # df['cause_persian'] = [disease.cause_persian for disease in data['diseases']]
+    
+    mlb = MultiLabelBinarizer(sparse_output=True)
+    
+    for col in df.columns:
+        if col == 'disease':
+            continue
+        
+        try:
+            df = df.join(
+                pd.DataFrame.sparse.from_spmatrix(
+                    mlb.fit_transform(df.pop(col)), # type: ignore
+                    index=df.index,
+                    columns=[f'{col}_'] + mlb.classes_
+                ),
+            )
+        except Exception as error:
+            print(f'Error: {error} at column: {col}, skipping...')
+        
+    
+    return df
+
+
+def trainingDataFromPromptsForBERT(data: dict) -> pd.DataFrame:
+    """
+    Get the training data from the prompts JSON file
+    ================================================
+    Parameters:
+    -----------
+        data:
+            description: Prompts JSON file
+            type: dict
+    -----------
+    Returns:
+    --------
+        data:
+            description: Training data
+            type: pd.DataFrame
+    """
+    
+    sentences = []
+    
+    for prompt in data['diseasesPrompts']:
+        for sentence in prompt['sentences']:
+            sentences.append((sentence, prompt['disease']))
+    
+    df = pd.DataFrame(sentences, columns=['sentence', 'disease'])
+    
+    return df

tasks/data/utility.py

@@ -0,0 +1,57 @@
+import pandas as pd
+
+
+def getSymptomsCausesAndDiseaseNameFromJSON(data: pd.DataFrame) -> dict:
+    """
+    Generate Symptoms, Causes and Disease Name
+    =========================
+    Parameters:
+    -----------
+        data:
+            description: Augmented Data
+            type: pd.DataFrame
+    -----------
+    Returns:
+    --------
+        data:
+            description: Symptoms, Causes and Disease Name
+            type: dict
+    --------------------------------------------------------------------------------------------
+    Working:
+    --------
+        - Create a DataFrame from data
+        - Create a new DataFrame with columns: disease, symptoms and causes
+        - For each row in data:
+            - Add all the column names that are 1 in array to the 'symptoms' key if the prefix is symptom_ otherwise to 'causes' key and set the value of 'disease' key to disease name
+        - Return the new dict
+    """
+    
+    
+    numberOfSymptoms = len([col for col in data.columns if col.startswith('symptoms_')])
+    
+    diseases = data['disease']
+    symptoms = data.columns[1:numberOfSymptoms]
+    causes = data.columns[numberOfSymptoms:]   
+    
+    # Get symptoms column names where symptoms = 1
+    symptomsArray = data[symptoms].apply(lambda x: symptoms[x.values.astype(bool)].tolist(), axis=1)
+    
+    # Get causes column names where causes = 1
+    causesArray = data[causes].apply(lambda x: causes[x.values.astype(bool)].tolist(), axis=1)
+    
+    # Remove the prefix 'symptoms_' and 'causes_' from the symptomsArray and causesArray
+    symptomsArray = [list(map(lambda x: x.replace('symptoms_', ''), symptom)) for symptom in symptomsArray]
+    causesArray = [list(map(lambda x: x.replace('causes_', ''), cause)) for cause in causesArray]
+    
+    diseaseDict = {
+        "diseases": [
+            {
+                "disease": disease,
+                "symptoms": symptom,
+                "causes": cause
+            }
+            for disease, symptom, cause in zip(diseases, symptomsArray, causesArray)
+        ]
+    }
+    
+    return diseaseDict

tasks/training/handle_train.py

@@ -0,0 +1,84 @@
+import argparse
+import os
+
+import boto3
+import datasets
+import sagemaker
+import sagemaker.s3 as S3Downloader
+import transformers
+from sagemaker.huggingface import HuggingFace
+
+
+def train(trainDataPath: str, testDataPath: str, hyperparameters: dict | None = None):
+    sess = sagemaker.Session()
+    # sagemaker session bucket -> used for uploading data, models and logs
+    # sagemaker will automatically create this bucket if it not exists
+    sagemaker_session_bucket=None
+    if sagemaker_session_bucket is None and sess is not None:
+        # set to default bucket if a bucket name is not given
+        sagemaker_session_bucket = sess.default_bucket()
+
+    try:
+        role = sagemaker.get_execution_role()
+    except ValueError:
+        iam = boto3.client('iam')
+        role = iam.get_role(RoleName='sagemaker_execution_role')['Role']['Arn']
+
+    sess = sagemaker.Session(default_bucket=sagemaker_session_bucket)
+    
+    
+    tokenizer_name = 'HaiderSultanArc/UnaniBERT'
+    tokenizer = transformers.BertTokenizer.from_pretrained(tokenizer_name)
+    
+    def tokenize(batch):
+        return tokenizer(batch['sentence'], padding='max_length', truncation=True)
+    
+    
+    train_dataset = datasets.load_from_disk(trainDataPath)
+    test_dataset = datasets.load_from_disk(testDataPath)
+    
+    train_dataset = train_dataset.map(tokenize, batched=True)
+    test_dataset = test_dataset.map(tokenize, batched=True)
+    
+    train_dataset = train_dataset.rename_column("disease", "labels")
+    test_dataset = test_dataset.rename_column("disease", "labels")
+    
+    train_dataset.set_format('torch', columns=['input_ids', 'attention_mask', 'labels'])
+    test_dataset.set_format('torch', columns=['input_ids', 'attention_mask', 'labels'])
+    
+    # save train_dataset to s3
+    training_input_path = f's3://{sess.default_bucket()}/UnaniBERT_dataset/train'
+    train_dataset.save_to_disk(training_input_path)
+
+    # save test_dataset to s3
+    test_input_path = f's3://{sess.default_bucket()}/UnaniBERT_dataset/test'
+    test_dataset.save_to_disk(test_input_path)
+
+    # hyperparameters which are passed to the training job
+    hyperparameters = {
+        'epochs': 50,
+        'train_batch_size': 32,
+        'model_name': 'HaiderSultanArc/UnaniBERT'
+    } if hyperparameters is None else hyperparameters
+
+    # create the Estimator
+    huggingface_estimator = HuggingFace(
+        entry_point='train.py',
+        source_dir='./tasks/training',
+        instance_type='ml.p3.2xlarge',
+        instance_count=1,
+        role=role,
+        transformers_version='4.26',
+        pytorch_version='1.13',
+        py_version='py39',
+        hyperparameters = hyperparameters
+    )
+    
+    huggingface_estimator.fit(
+        {
+            'train': trainDataPath,
+            'test': testDataPath
+        }
+    )
+    
+    predictor = huggingface_estimator.deploy(1, "ml.g4dn.xlarge")

tasks/training/helper.py

@@ -0,0 +1,47 @@
+import datetime
+import time
+
+import matplotlib.pyplot as plt
+import numpy as np
+import seaborn as sns
+import torch
+
+
+def getDevice():
+    if torch.cuda.is_available():    
+        device = torch.device("cuda")
+
+        print('There are %d GPU(s) available.' % torch.cuda.device_count())
+        print('We will use the GPU:', torch.cuda.get_device_name(0))
+    else:
+        print('No GPU available, using the CPU instead.')
+        device = torch.device("cpu")
+    
+    return device
+
+
+def flatAccuracy(preds, labels):
+        pred_flat = np.argmax(preds, axis=1).flatten()
+        labels_flat = labels.flatten()
+        return np.sum(pred_flat == labels_flat) / len(labels_flat)
+    
+
+def formatTime(elapsed):
+    elapsed_rounded = int(round((elapsed)))
+    
+    # Format as hh:mm:ss
+    return str(datetime.timedelta(seconds=elapsed_rounded))
+
+
+def plotTrainingLoss(lossValues):
+    sns.set(style='darkgrid')
+    sns.set(font_scale=1.5)
+    plt.rcParams["figure.figsize"] = (12,6)
+
+    plt.plot(lossValues, 'b-o')
+
+    plt.title("Training loss")
+    plt.xlabel("Epoch")
+    plt.ylabel("Loss")
+
+    plt.show()

tasks/training/preprocessing.py

@@ -0,0 +1,109 @@
+import pandas as pd
+import torch
+import transformers
+from sklearn import preprocessing
+from sklearn.model_selection import train_test_split
+from torch.nn import ConstantPad1d
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import (DataLoader, RandomSampler, SequentialSampler,
+                              TensorDataset)
+
+
+def preprocess(data: pd.DataFrame, tokenizer: transformers.BertTokenizer):
+    """
+    Preprocesses the data.
+    ======================
+    Parameters:
+        data (pd.DataFrame): The data to be preprocessed.
+    ----------------------
+    Returns:
+        trainDataloader (DataLoader): The training dataloader.
+        validationDataloader (DataLoader): The validation dataloader.
+        testDataloader (DataLoader): The testing dataloader.
+        labelEncoder (LabelEncoder): The label encoder.
+    """
+    sentences = data.sentence.values
+    labels = data.disease.values
+    
+    encodedSentences = []
+
+    for sentence in sentences:
+        # `encode` will:
+        #   (1) Tokenize the sentence.
+        #   (2) Prepend the `[CLS]` token to the start.
+        #   (3) Append the `[SEP]` token to the end.
+        #   (4) Map tokens to their IDs.
+        encSentence = tokenizer.encode(
+                            sentence,                      # Sentence to encode.
+                            add_special_tokens = True, # Add '[CLS]' and '[SEP]'
+
+                            # This function also supports truncation and conversion
+                            # to pytorch tensors, but we need to do padding, so we
+                            # can't use these features :( .
+                            #max_length = 128,          # Truncate all sentences.
+                            #return_tensors = 'pt',     # Return pytorch tensors.
+                    )
+        
+        # Add the encoded sentence to the list.
+        encodedSentences.append(encSentence)
+    
+    MAX_LEN = max([len(sen) for sen in encodedSentences]) + 10
+    
+    seq = [torch.tensor(sen) for sen in encodedSentences]
+    padSequences = [ConstantPad1d((0, MAX_LEN - len(sen)), 0)(sen) for sen in seq]
+    encodedSentences = pad_sequence(padSequences, batch_first=True)
+    
+    attentionMasks = []
+
+    for sentence in encodedSentences:
+        # Create the attention mask.
+        #   - If a token ID is 0, then it's padding, set the mask to 0.
+        #   - If a token ID is > 0, then it's a real token, set the mask to 1.
+        attMask = [int(token_id > 0) for token_id in sentence]
+        
+        # Store the attention mask for this sentence.
+        attentionMasks.append(attMask)
+    
+    
+    labelEncoder = preprocessing.LabelEncoder()
+    labels = labelEncoder.fit_transform(labels) # type: ignore
+
+    trainingSentences, testingSentences, trainingLabels, testingLabels = train_test_split(encodedSentences, labels, test_size=0.3, random_state=2018)
+
+    # Use 90% for training and 10% for validation.
+    trainInputs, validationInputs, trainLabels, validationLabels = train_test_split(trainingSentences, trainingLabels, 
+                                                                random_state=2018, test_size=0.1)
+    # Do the same for the masks.
+    trainingMasks, testingMasks, _, _ = train_test_split(attentionMasks, labels, random_state=2018, test_size=0.3)
+
+    train_masks, validationMasks, _, _ = train_test_split(trainingMasks, trainingLabels, random_state=2018, test_size=0.1)
+    
+    # Convert all inputs and labels into torch tensors, the required datatype 
+    # for our model.
+    trainInputs = torch.tensor(trainInputs)
+    validationInputs = torch.tensor(validationInputs)
+    testInputs = torch.tensor(testingSentences)
+
+    trainLabels = torch.tensor(trainLabels)
+    validationLabels = torch.tensor(validationLabels)
+    testLabels = torch.tensor(testingLabels)
+
+    train_masks = torch.tensor(train_masks)
+    validationMasks = torch.tensor(validationMasks)
+    testMasks = torch.tensor(testingMasks)
+
+    batchSize = 32
+
+    trainData = TensorDataset(trainInputs, train_masks, trainLabels)
+    trainSampler = RandomSampler(trainData)
+    trainDataloader = DataLoader(trainData, sampler=trainSampler, batch_size=batchSize)
+
+    validationData = TensorDataset(validationInputs, validationMasks, validationLabels)
+    validationSampler = SequentialSampler(validationData)
+    validationDataloader = DataLoader(validationData, sampler=validationSampler, batch_size=batchSize)
+    
+    testData = TensorDataset(testInputs, testMasks, testLabels)
+    testSampler = SequentialSampler(testData)
+    testDataloader = DataLoader(testData, sampler=testSampler, batch_size=batchSize)
+    
+    return trainDataloader, validationDataloader, testDataloader, labelEncoder

tasks/training/train.py

@@ -0,0 +1,95 @@
+import argparse
+import logging
+import os
+import sys
+
+import transformers
+from datasets import load_from_disk
+from sklearn.metrics import accuracy_score, precision_recall_fscore_support
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+
+    # hyperparameters sent by the client are passed as command-line arguments to the script.
+    parser.add_argument("--epochs", type=int, default=3)
+    parser.add_argument("--train_batch_size", type=int, default=32)
+    parser.add_argument("--eval_batch_size", type=int, default=64)
+    parser.add_argument("--warmup_steps", type=int, default=500)
+    parser.add_argument("--model_name", type=str)
+    parser.add_argument("--learning_rate", type=str, default=5e-5)
+
+    # Data, model, and output directories
+    parser.add_argument("--output_data_dir", type=str, default=os.environ["SM_OUTPUT_DATA_DIR"])
+    parser.add_argument("--model_dir", type=str, default=os.environ["SM_MODEL_DIR"])
+    parser.add_argument("--n_gpus", type=str, default=os.environ["SM_NUM_GPUS"])
+    parser.add_argument("--training_dir", type=str, default=os.environ["SM_CHANNEL_TRAIN"])
+    parser.add_argument("--test_dir", type=str, default=os.environ["SM_CHANNEL_TEST"])
+    
+    args, _ = parser.parse_known_args()
+
+    # Set up logging
+    logger = logging.getLogger(__name__)
+
+    logging.basicConfig(
+        level=logging.getLevelName("INFO"),
+        handlers=[logging.StreamHandler(sys.stdout)],
+        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    )
+    
+     # load datasets
+    train_dataset = load_from_disk(args.training_dir)
+    test_dataset = load_from_disk(args.test_dir)
+
+    logger.info(f" loaded train_dataset length is: {len(train_dataset)}")
+    logger.info(f" loaded test_dataset length is: {len(test_dataset)}")
+
+    # compute metrics function for binary classification
+    def compute_metrics(pred):
+        labels = pred.label_ids
+        preds = pred.predictions.argmax(-1)
+        precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average="binary")
+        acc = accuracy_score(labels, preds)
+        return {"accuracy": acc, "f1": f1, "precision": precision, "recall": recall}
+    
+    
+    # download model from model hub
+    model = transformers.BertForSequenceClassification.from_pretrained(args.model_name)
+    tokenizer = transformers.BertTokenizer.from_pretrained(args.model_name)
+    
+    # define training args
+    training_args = transformers.TrainingArguments(
+        output_dir=args.model_dir,
+        num_train_epochs=args.epochs,
+        per_device_train_batch_size=args.train_batch_size,
+        per_device_eval_batch_size=args.eval_batch_size,
+        warmup_steps=args.warmup_steps,
+        evaluation_strategy="epoch",
+        logging_dir=f"{args.output_data_dir}/logs",
+        learning_rate=float(args.learning_rate),
+    )
+
+    # create Trainer instance
+    trainer = transformers.Trainer(
+        model=model,
+        args=training_args,
+        compute_metrics=compute_metrics,
+        train_dataset=train_dataset,
+        eval_dataset=test_dataset,
+        tokenizer=tokenizer,
+    )
+
+    # train model
+    trainer.train()
+
+    # evaluate model
+    eval_result = trainer.evaluate(eval_dataset=test_dataset)
+
+    # writes eval result to file which can be accessed later in s3 ouput
+    with open(os.path.join(args.output_data_dir, "eval_results.txt"), "w") as writer:
+        print(f"***** Eval results *****")
+        for key, value in sorted(eval_result.items()):
+            writer.write(f"{key} = {value}\n")
+
+    # Saves the model to s3
+    trainer.save_model(args.model_dir)

utils/handle.py

@@ -0,0 +1,114 @@
+import json
+import os
+
+import pandas as pd
+from fastapi import UploadFile
+
+import tasks.data.dataAugmentation as da
+import tasks.data.dataEngineering as de
+import tasks.data.utility as util
+import tasks.training.handle_train as trainingPipeline
+
+
+def augmentDataUsingVectorSpaceAlgorithm(file: UploadFile, savePath: str):
+    try:
+        os.makedirs(os.path.dirname(savePath), exist_ok=True)
+        data = pd.read_csv(file.file)
+        df = da.augmentDataWithVectorSpaceAlgorithm(data)
+        df.to_csv(savePath, index=False, encoding='utf-8')
+        
+        return {
+            "success": True,
+            "message": "Training data augmented successfully",
+            "data": df.head(5).to_dict()
+        }
+    except Exception as error:
+        return {
+            "success": False,
+            "message": f"Training data augmentation failed. {error}",
+            "data": None
+        }
+
+
+def getSymptomsCausesAndDiseaseNameFromJSON(file: UploadFile, savePath: str):
+    try:
+        os.makedirs(os.path.dirname(savePath), exist_ok=True)
+        data = pd.read_csv(file.file)
+        
+        diseaseDict = util.getSymptomsCausesAndDiseaseNameFromJSON(data)
+        
+        json.dump(diseaseDict, open(savePath, 'w', encoding='utf-8'), ensure_ascii=False)
+        
+        return {
+            "success": True,
+            "message": "Symptoms, causes and disease name extracted successfully",
+            "data": None
+        }
+    except Exception as error:
+        return {
+            "success": False,
+            "message": f"Symptoms, causes and disease name extraction failed. {error}",
+            "data": None
+        }
+
+
+def trainingDataFromUTagsJSON(file: UploadFile, savePath: str):
+    try:
+        os.makedirs(os.path.dirname(savePath), exist_ok=True)
+        data = json.loads(file.file.read())
+        df = de.trainingDataFromUTagsJSON(data)
+        df.to_csv(savePath, index=False, encoding='utf-8')
+        
+        return {
+            "success": True,
+            "message": "Training data generated successfully",
+            "data": df.head(5).to_dict()
+        }
+    except Exception as error:
+        return {
+            "success": False,
+            "message": f"Training data generation failed. {error}",
+            "data": None
+        }
+
+
+def trainingDataFromPromptsForBERT(file: UploadFile, savePath: str):
+    try:
+        os.makedirs(os.path.dirname(savePath), exist_ok=True)
+        data = json.loads(file.file.read())
+        df = de.trainingDataFromPromptsForBERT(data)
+        df.to_csv(savePath, index=False, encoding='utf-8')
+        
+        return {
+            "success": True,
+            "message": "Training data generated successfully",
+            "data": df.head(5).to_dict()
+        }
+    except Exception as error:
+        return {
+            "success": False,
+            "message": f"Training data generation failed. {error}",
+            "data": None
+        }
+    
+
+def trainModelOnSageMaker(trainDataPath: str, testDataPath: str, file: UploadFile | None = None):
+    try:
+        hyperparameters = None
+        if file is not None:
+            hyperparameters = json.loads(file.file.read())
+        
+        trainingPipeline.train(trainDataPath, testDataPath, hyperparameters)
+        
+        return {
+            "success": True,
+            "message": "Model trained successfully",
+            "data": None
+        }
+    except Exception as error:
+        return {
+            "success": False,
+            "message": f"Model training failed. {error}",
+            "data": None
+        }
+