First commit

Dockerfile

@@ -0,0 +1,21 @@
+# This Dockerfile serves as the build file for Huggingface Spaces
+FROM huggingface/transformers-pytorch-gpu
+
+ARG ML_APP_LISTEN_PORT=7860
+ARG ML_MLFLOW_ENDPOINT
+ARG ML_HF_ACCESS_TOKEN
+
+ENV APP_LISTEN_PORT=${ML_APP_LISTEN_PORT}
+ENV MLFLOW_ENDPOINT=${ML_MLFLOW_ENDPOINT}
+ENV HF_ACCESS_TOKEN=${ML_HF_ACCESS_TOKEN}
+
+RUN apt-get update
+RUN /usr/bin/python3 -m pip install uvicorn fastapi mlflow huggingface_hub httpx
+
+WORKDIR /app
+
+COPY ./app /app
+
+EXPOSE ${ML_APP_LISTEN_PORT}
+
+ENTRYPOINT [ "/usr/bin/python3", "/app/main.py" ]

README.md

@@ -0,0 +1,8 @@
+---
+title: SocialMedia Sentiment Analysis
+emoji: 🐳
+colorFrom: purple
+colorTo: gray
+sdk: docker
+app_port: 7860
+---

app/config.py

@@ -0,0 +1,13 @@
+import os
+
+def is_test_mode():
+	"""
+	Checks whether test mode is enabled or not.
+
+	Returns:
+		bool: True if test mode is enabled, false otherwise
+	"""
+	return os.environ.get("TEST_MODE", None) == "1"
+
+def enable_test_mode():
+    os.environ["TEST_MODE"] = "1"

app/main.py

@@ -0,0 +1,115 @@
+from fastapi import FastAPI, BackgroundTasks
+from fastapi.responses import JSONResponse, HTMLResponse
+from pydantic import BaseModel
+import uvicorn
+from uvicorn.config import logger
+import os
+import argparse
+from tasks.training import TrainingTask
+from config import enable_test_mode
+
+app = FastAPI()
+
+@app.post("/train/start", response_class=JSONResponse)
+async def start_model_training(background_tasks: BackgroundTasks):
+	"""
+	Endpoint on which a request can be sent to start model re-training, 
+	if there's no training task currently running.
+	The task will be carried out in background and its status can be
+	polled via /train/get_state.
+
+	Args:
+		background_tasks (BackgroundTasks): BG Tasks scheduler provided by FastAPI
+
+	Returns:
+		dict: A dictionary containing a message of the outcome for the request.
+	"""
+	if not TrainingTask.has_instance():
+		background_tasks.add_task(TrainingTask.get_instance())
+
+		return {
+			"message": "Model training was scheduled and will begin shortly.",
+		}
+	
+	return {
+		"message": "A training instance is already running.",
+	}
+
+@app.post("/train/get_state", response_class=JSONResponse)
+async def poll_model_training_state():
+	"""
+	Checks if there is currently a training task ongoing.
+	If so, returns whether it's done and/or if an error occurred.
+	Otherwise if no instance is running, returns only a message.
+
+	Returns:
+		dict: Dictionary containing either done/error or message.
+	"""
+	if TrainingTask.has_instance():
+		train_instance : TrainingTask = TrainingTask.get_instance()
+		is_done = train_instance.is_done()
+		has_error = train_instance.has_error()
+
+		if is_done:
+			TrainingTask.clear_instance()
+
+		return {
+			"done": is_done,
+			"error": has_error,
+		}
+
+	return {
+		"message": "No training instance running!",
+	}
+
+class InferenceRequest(BaseModel):
+	"""
+	Provides a model/schema for the accepted request body for incoming
+	inference requests.
+	"""
+	messages: list[str]
+
+@app.post("/inference", response_class=JSONResponse)
+async def inference(data: InferenceRequest):
+	"""
+	
+
+	Args:
+		data (InferenceRequest): Structure containing a list of 
+  								 messages that shall be evaluated
+
+	Returns:
+		json: A json list containing the sentiment analysis for each message.
+			  Each element consists of a dictionary with the following keys:
+			  positive, neutral, negative
+	"""
+	
+	from tasks.inference import infer_task
+	return infer_task.predict(data.messages)
+
+@app.get("/", response_class=HTMLResponse)
+async def root():
+	"""
+	The root endpoint for our hosted application. Only shows a message
+	showing that it's up and running.
+
+	Returns:
+		str: A html response containing a hello world-like string
+	"""
+	return "Hi there! It's a nice blank page, isn't it?"
+
+if __name__ == "__main__":
+	"""
+	Entrypoint for the application executed via command-line.
+	It accepts an optional argument "--test" to enable the test mode.
+	"""
+	parser = argparse.ArgumentParser()
+	parser.add_argument("test", nargs="?", default="no")
+	args = parser.parse_args()
+
+	if args.test == "yes":
+		enable_test_mode()
+
+	config = uvicorn.Config("main:app", host="0.0.0.0", port=int(os.environ["APP_LISTEN_PORT"]), log_level="debug")
+	server = uvicorn.Server(config)
+	server.run()

app/tasks/inference.py

@@ -0,0 +1,115 @@
+from uvicorn.config import logger
+from transformers import AutoModelForSequenceClassification
+from transformers import AutoTokenizer, AutoConfig
+import numpy as np
+import mlflow
+import os
+import time
+from scipy.special import softmax
+
+# HuggingFace Model to be used for inferencing
+MODEL = f"cardiffnlp/twitter-roberta-base-sentiment-latest"
+
+class InferenceTask:
+
+	def __init__(self):
+		self.clear()
+		self.load_model()
+
+	def load_model(self):
+		try:
+			self.__tokenizer = AutoTokenizer.from_pretrained(MODEL)
+			self.__config = AutoConfig.from_pretrained(MODEL)
+			self.__model = AutoModelForSequenceClassification.from_pretrained(MODEL)
+			self.__is_loaded = True
+		except Exception as ex:
+			logger.error("Failed to load inference model: {ex}")
+			self.clear()
+			return False
+
+		return True
+
+	def clear(self):
+		self.__is_loaded = False
+		self.__tokenizer = None
+		self.__config = None
+		self.__model = None
+
+	def is_loaded(self):
+		return self.__is_loaded
+
+	def predict(self, messages: list[str]):
+		if len(messages) == 0:
+			return None
+
+		if not self.is_loaded() and not self.load_model():
+			return None
+
+		mlflow.set_tracking_uri(os.environ["MLFLOW_ENDPOINT"])
+		mlflow.set_experiment("Sentiment Analysis")
+
+		with mlflow.start_run() as run:
+			preprocessed_messages = self.__preprocess(messages)
+			labelized_scores = []
+
+			for message in preprocessed_messages:
+				encoded_input = self.__tokenizer(message, return_tensors='pt', padding="longest")
+				output = self.__model(**encoded_input)
+				scores = output[0][0].detach().numpy()
+				scores = softmax(scores)
+				scores = self.__labelize(scores)
+				labelized_scores.append(scores)
+
+			mean_sentiment = self.__calculate_mean_sentiment(labelized_scores)
+			mean_sentiment["samples"] = len(labelized_scores)
+			logger.info(mean_sentiment)
+
+			mlflow.log_metrics(mean_sentiment, step=int(time.time()))
+
+			return labelized_scores
+
+	def __calculate_mean_sentiment(self, labelized_scores: list):
+		total_samples = float(len(labelized_scores))
+
+		mean_sentiment = {
+			"positive": 0.0,
+			"neutral": 0.0,
+			"negative": 0.0,
+		}
+
+		for score in labelized_scores:
+			mean_sentiment["positive"] += score["positive"]
+			mean_sentiment["neutral"] += score["neutral"]
+			mean_sentiment["negative"] += score["negative"]
+
+		mean_sentiment["positive"] /= total_samples
+		mean_sentiment["neutral"] /= total_samples
+		mean_sentiment["negative"] /= total_samples
+
+		return mean_sentiment
+
+	# Preprocess text (username and link placeholders)
+	def __preprocess(self, messages: list[str]):
+		msg_list = []
+		for message in messages:
+			new_message = []
+			for t in message.split(" "):
+				t = '@user' if t.startswith('@') and len(t) > 1 else t
+				t = 'http' if t.startswith('http') else t
+				new_message.append(t)
+			msg_list.append(" ".join(new_message))
+		return msg_list
+
+	def __labelize(self, scores):
+		output = {}
+		ranking = np.argsort(scores)
+		ranking = ranking[::-1]
+		for i in range(scores.shape[0]):
+			l = self.__config.id2label[ranking[i]]
+			s = float(scores[ranking[i]])
+			output[l] = s
+		return output
+
+# Preload a global instance so that inference can be
+# executed immediately when requested
+infer_task = InferenceTask()

app/tasks/training.py

@@ -0,0 +1,246 @@
+import evaluate
+import numpy as np
+from uvicorn.config import logger
+from datasets import load_dataset
+from transformers import (
+	AutoModelForSequenceClassification,
+	AutoTokenizer,
+	Trainer,
+	TrainingArguments,
+	pipeline,
+)
+from huggingface_hub import login, logout
+
+import os
+import mlflow
+from tasks.inference import infer_task
+from config import is_test_mode
+
+"""
+Documentation:
+- https://huggingface.co/docs/transformers/en//training
+- https://mlflow.org/docs/latest/llms/transformers/tutorials/fine-tuning/transformers-fine-tuning
+"""
+
+MODEL = "cardiffnlp/twitter-roberta-base-sentiment-latest"
+DATASET = "zeroshot/twitter-financial-news-sentiment"
+HF_DEST_REPO = "financial-twitter-roberta-sentiment"
+
+RNG_SEED = 22
+
+class TrainingTask:
+
+	TRAINING_TASK_INST_SINGLETON = None
+	
+	def __init__(self):
+		self.__is_done = False
+		self.__has_error = False
+
+		self.__train_dataset = None
+		self.__test_dataset = None
+		self.__tokenizer = None
+		self.__train_tokenized = None
+		self.__test_tokenized = None
+		self.__model = None
+		self.__trainer = None
+		self.__run_id = None
+
+	@staticmethod
+	def has_instance():
+		return TrainingTask.TRAINING_TASK_INST_SINGLETON is not None
+
+	@staticmethod
+	def get_instance():
+		if TrainingTask.TRAINING_TASK_INST_SINGLETON is None:
+			TrainingTask.TRAINING_TASK_INST_SINGLETON = TrainingTask()
+		
+		return TrainingTask.TRAINING_TASK_INST_SINGLETON
+
+	@staticmethod
+	def clear_instance():
+		del TrainingTask.TRAINING_TASK_INST_SINGLETON
+		TrainingTask.TRAINING_TASK_INST_SINGLETON = None
+
+	def has_error(self):
+		return self.__has_error
+
+	def is_done(self):
+		return self.__is_done
+
+	def __call__(self, *args, **kwds):
+		self.__has_error = False
+		self.__is_done = False
+
+		if is_test_mode():
+			# Simulate a successful training run in test mode
+			self.__has_error = False
+			self.__is_done = True
+			return
+
+		login(token=os.environ["HF_ACCESS_TOKEN"])
+
+		try:
+			self.__load_datasets()
+			self.__tokenize()
+			self.__load_model()
+			self.__train()
+			self.__evaluate()
+			self.__deploy()
+		except Exception as ex:
+			logger.error(f"Error during training: {ex}")
+			self.__has_error = True
+		finally:
+			self.__is_done = True
+
+		logout()
+
+		self.__reload_inference_model()
+
+	def __load_datasets(self):
+		# Load the dataset.
+		dataset = load_dataset(DATASET)
+
+		# Split train/test by an 8/2 ratio.
+		dataset_train_test = dataset["train"].train_test_split(test_size=0.2)
+		self.__train_dataset = dataset_train_test["train"]
+		self.__test_dataset = dataset_train_test["test"]
+
+		# Swap labels so that they match what the model actually expects
+		# The model expects {0: positive, 1: neutral, 2: negative}
+		# But the dataset uses {0: positive, 1: negative, 2: neutral}
+		# So here we just flip 1<->2 to remain consistent
+		def label_filter(row):
+			row["label"] = { 0: 0, 1: 2, 2: 1 }[row["label"]]
+			return row
+
+		self.__train_dataset = self.__train_dataset.map(label_filter)
+		self.__test_dataset = self.__test_dataset.map(label_filter)
+
+	def __tokenize(self):
+		# Load the tokenizer for the model.
+		self.__tokenizer = AutoTokenizer.from_pretrained(MODEL)
+
+		def tokenize_function(examples):
+			# Pad/truncate each text to 512 tokens. Enforcing the same shape
+			# could make the training faster.
+			return self.__tokenizer(
+				examples["text"],
+				padding="max_length",
+				truncation=True,
+				max_length=256,
+			)
+
+		# Tokenize the train and test datasets
+		self.__train_tokenized = self.__train_dataset.map(tokenize_function)
+		self.__train_tokenized = self.__train_tokenized.remove_columns(["text"]).shuffle(seed=RNG_SEED)
+
+		self.__test_tokenized = self.__test_dataset.map(tokenize_function)
+		self.__test_tokenized = self.__test_tokenized.remove_columns(["text"]).shuffle(seed=RNG_SEED)
+
+	def __load_model(self):
+		# Set the mapping between int label and its meaning.
+		id2label = {0: "Bearish", 1: "Neutral", 2: "Bullish"}
+		label2id = {"Bearish": 0, "Neutral": 1, "Bullish": 2}
+
+		# Acquire the model from the Hugging Face Hub, providing label and id mappings so that both we and the model can 'speak' the same language.
+		self.__model = AutoModelForSequenceClassification.from_pretrained(
+			MODEL,
+			num_labels=3,
+			label2id=label2id,
+			id2label=id2label,
+		)
+
+	def __train(self):
+		# Define the target optimization metric
+		metric = evaluate.load("accuracy")
+
+		# Define a function for calculating our defined target optimization metric during training
+		def compute_metrics(eval_pred):
+			logits, labels = eval_pred
+			predictions = np.argmax(logits, axis=-1)
+			return metric.compute(predictions=predictions, references=labels)
+
+		# Checkpoints will be output to this `training_output_dir`.
+		training_output_dir = "/tmp/sentiment_trainer"
+		training_args = TrainingArguments(
+			output_dir=training_output_dir,
+			eval_strategy="epoch",
+			per_device_train_batch_size=8,
+			per_device_eval_batch_size=8,
+			logging_steps=8,
+			num_train_epochs=3,
+		)
+
+		# Instantiate a `Trainer` instance that will be used to initiate a training run.
+		self.__trainer = Trainer(
+			model=self.__model,
+			args=training_args,
+			train_dataset=self.__train_tokenized,
+			eval_dataset=self.__test_tokenized,
+			compute_metrics=compute_metrics,
+		)
+
+		mlflow.set_tracking_uri(os.environ["MLFLOW_ENDPOINT"])
+		mlflow.set_experiment("Sentiment Classifier Training")
+
+		with mlflow.start_run() as run:
+			self.__run_id = run.info.run_id
+			self.__trainer.train()
+
+	def __evaluate(self):
+		tuned_pipeline = pipeline(
+			task="text-classification",
+			model=self.__trainer.model,
+			batch_size=8,
+			tokenizer=self.__tokenizer,
+			device="cpu", # or cuda
+		)
+
+		quick_check = (
+			"I have a question regarding the project development timeline and allocated resources; "
+			"specifically, how certain are you that John and Ringo can work together on writing this next song? "
+			"Do we need to get Paul involved here, or do you truly believe, as you said, 'nah, they got this'?"
+		)
+
+		result = tuned_pipeline(quick_check)
+		logger.debug("Test evaluation of fine-tuned model: %s %.6f" % (result[0]["label"], result[0]["score"]))
+
+		# Define a set of parameters that we would like to be able to flexibly override at inference time, along with their default values
+		model_config = {"batch_size": 8}
+
+		# Infer the model signature, including a representative input, the expected output, and the parameters that we would like to be able to override at inference time.
+		signature = mlflow.models.infer_signature(
+			["This is a test!", "And this is also a test."],
+			mlflow.transformers.generate_signature_output(
+				tuned_pipeline, ["This is a test response!", "So is this."]
+			),
+			params=model_config,
+		)
+		
+		# Log the pipeline to the existing training run
+		with mlflow.start_run(run_id=self.__run_id):
+			model_info = mlflow.transformers.log_model(
+				transformers_model=tuned_pipeline,
+				artifact_path="fine_tuned",
+				signature=signature,
+				input_example=["Pass in a string", "And have it mark as spam or not."],
+				model_config=model_config,
+			)
+
+			# Load our saved model in the native transformers format
+			loaded = mlflow.transformers.load_model(model_uri=model_info.model_uri)
+
+			# Define a test example that we expect to be classified as spam
+			validation_text = (
+				"Want to learn how to make MILLIONS with no effort? Click HERE now! See for yourself! Guaranteed to make you instantly rich! "
+				"Don't miss out you could be a winner!"
+			)
+
+			# validate the performance of our fine-tuning
+			loaded(validation_text)
+
+	def __deploy(self):
+		self.__trainer.push_to_hub(HF_DEST_REPO)
+
+	def __reload_inference_model(self):
+		infer_task.load_model()

app/test/fixture.py

@@ -0,0 +1,19 @@
+import pytest
+from fastapi.testclient import TestClient
+from main import app
+import os
+
+@pytest.fixture()
+def app_client():
+	"""
+	Barebone test fixture that initializes a FastAPI TestClient
+	which can be used to test all endpoints provided by the application.
+
+	Yields:
+		TestClient: A client hosting the whole application so that it
+					can be accessed and controlled programmatically.
+	"""
+	os.environ["TEST_MODE"] = "1" # Turns off actual model training
+	
+	client = TestClient(app)
+	yield client

app/test/test_inferencing_endpoint.py

@@ -0,0 +1,86 @@
+from .fixture import app_client
+import json
+
+def test_inference_endpoint(app_client):
+	"""
+	Tests the output of the inference endpoint of the application.
+	
+	Given:
+		- Payload with valid list of messages
+	When:
+		- POST Request sent to inference endpoint
+	Then:
+		- Expect message to be classified as positive
+	"""
+	response = app_client.post(
+		"/inference",
+		headers={
+			"Content-Type": "application/json",
+   		},
+		json={
+			"messages": [
+				"BTC is going to skyrocket!",
+			],
+		}
+	)
+
+	assert response.status_code == 200
+	output = json.loads(response.content)
+	assert isinstance(output, list)
+	assert len(output) == 1
+	assert output[0]["positive"] > output[0]["negative"] and output[0]["positive"] > output[0]["neutral"]
+
+
+def test_inference_endpoint_with_wrong_payload(app_client):
+	"""
+	Tests the output of the inference endpoint of the application with an
+	invalid payload.
+	This should yield a 422 status error as FastAPI will not be able
+	to translate the payload into the InferenceRequest model.
+	
+	Given:
+		- Payload with wrong message key
+	When:
+		- POST Request sent to inference endpoint
+	Then:
+		- Expect 422 status code
+	"""
+	response = app_client.post(
+		"/inference",
+		headers={
+			"Content-Type": "application/json",
+   		},
+		json={
+			"msgs": [
+				"BTC is going to skyrocket!",
+			],
+		}
+	)
+
+	assert response.status_code == 422 # Unprocessable entity
+
+def test_inference_endpoint_with_no_prompt(app_client):
+	"""
+	Tests the output of the inference endpoint of the application
+	when a valid payload is provided but with no actual messages.
+	
+	Given:
+		- Payload without any messages
+	When:
+		- POST Request sent to inference endpoint
+	Then:
+		- Expect no error and correct format
+	"""
+	response = app_client.post(
+		"/inference",
+		headers={
+			"Content-Type": "application/json",
+   		},
+		json={
+			"messages": [],
+		}
+	)
+
+	assert response.status_code == 200
+	output = json.loads(response.content)
+	assert output == None

app/test/test_training_endpoint.py

@@ -0,0 +1,40 @@
+from .fixture import app_client
+import json
+import time
+
+def test_training_endpoint(app_client):
+	"""
+	Checks whether the training endpoint correctly receives, starts
+	and clears a training task instance.
+
+	Given:
+	- Launched training instance
+
+	When:
+	- State polled multiple times
+
+	Then:
+	- Expect state returned on first poll and instance gone on second poll
+	"""
+	response = app_client.post("/train/start")
+	assert response.status_code == 200
+	output : dict = json.loads(response.content)
+	assert len(output.keys()) == 1
+	assert output["message"] == "Model training was scheduled and will begin shortly."
+
+	time.sleep(5)
+
+	response = app_client.post("/train/get_state")
+	assert response.status_code == 200
+	output : dict = json.loads(response.content)
+	assert len(output.keys()) == 2
+	assert output["done"]
+	assert not output["error"]
+
+	time.sleep(1)
+
+	response = app_client.post("/train/get_state")
+	assert response.status_code == 200
+	output : dict = json.loads(response.content)
+	assert len(output.keys()) == 1
+	assert output["message"] == "No training instance running!"

docker-compose.test.yaml

@@ -0,0 +1,5 @@
+services:
+  # Override parameters of this service
+  model_runner:
+    # Sets the entrypoint so that pytest is executed
+    entrypoint: ["pytest"]

docker-compose.yaml

@@ -0,0 +1,25 @@
+services:
+  # 
+  model_runner:
+    build: 
+      context: .
+      dockerfile: Dockerfile
+
+    environment:
+      - APP_LISTEN_PORT=${APP_LISTEN_PORT}
+      - MLFLOW_ENDPOINT=${MLFLOW_ENDPOINT}
+      - HF_ACCESS_TOKEN=${HF_ACCESS_TOKEN}
+
+    container_name: mlrunner
+    restart: on-failure
+    ports:
+      - "${APP_LISTEN_PORT}:${APP_LISTEN_PORT}"
+    volumes:
+      - "./app:/app"
+    entrypoint: ["/usr/bin/python3", "/app/main.py"]
+    networks:
+      - airflow_tracking_network
+
+networks:
+  airflow_tracking_network:
+    external: true