Utilizing a Custom Model
While the process of building and fine-tuning a model is not the intent of this project, it is important to understand how a model can be added to this process.
# Filename: train.pyimport torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from datasets import load_dataset
from torch.utils.data import DataLoader
def train_model():
# Load dataset
full_dataset = load_dataset("stanfordnlp/imdb", split="train")
dataset = full_dataset.shuffle(seed=42).select(range(10000))
model_name = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
# Use GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.train()
# Create a DataLoader for batching
dataloader = DataLoader(dataset, batch_size=8, shuffle=True)
# Training loop
num_epochs = 3 # Set the number of epochs
for epoch in range(num_epochs):
total_loss = 0
for batch in dataloader:
inputs = tokenizer(batch["text"], truncation=True, padding=True, return_tensors="pt", max_length=512).to(device)
labels = torch.tensor(batch["label"]).to(device)
optimizer.zero_grad()
outputs = model(**inputs, labels=labels)
loss = outputs.loss
loss.backward()
optimizer.step()
total_loss += loss.item()
avg_loss = total_loss / len(dataloader)
print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {avg_loss:.4f}")
# Save the model
model.save_pretrained("./model/")
tokenizer.save_pretrained("./model/")
# Test the model with sample sentences
test_sentences = [
"This movie was fantastic!",
"I absolutely hated this film.",
"It was just okay, not great.",
"An absolute masterpiece!",
"Waste of time!",
"A beautiful story and well acted.",
"Not my type of movie.",
"It could have been better.",
"A thrilling adventure from start to finish!",
"Very disappointing."
]
# Switch model to evaluation mode
model.eval()
# Prepare tokenizer for test inputs
inputs = tokenizer(test_sentences, truncation=True, padding=True, return_tensors="pt", max_length=512).to(device)
with torch.no_grad():
outputs = model(**inputs)
predictions = torch.argmax(outputs.logits, dim=1)
# Print predictions
for sentence, prediction in zip(test_sentences, predictions):
sentiment = "positive" if prediction.item() == 1 else "negative"
print(f"Input: \"{sentence}\" -> Predicted sentiment: {sentiment}")
# Call the function to train the model and test it
train_model()
To make sure that we can query our new model that we have trained we have to update a few of our existing files. For instance, in main.py we now use the model from ./model and load it as a pretrained model. Additionally, for comparison’s sake, we add now have two endpoints to use, /predict/naive and predict/trained.
# Filename: main.pyfrom fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from transformers import pipeline
from prometheus_client import Counter, Histogram, start_http_server
import time
# Start prometheus metrics server on port 8001
start_http_server(8001)
app = FastAPI()
# Load the trained model and tokenizer from the local directory
model_path = "./model" # Path to your saved model
tokenizer = AutoTokenizer.from_pretrained(model_path)
trained_model = AutoModelForSequenceClassification.from_pretrained(model_path)
# Create pipelines
naive_classifier = pipeline("sentiment-analysis", device=-1)
trained_classifier = pipeline("sentiment-analysis", model=trained_model, tokenizer=tokenizer, device=-1)
# Metrics
PREDICTION_TIME = Histogram('prediction_duration_seconds', 'Time spent processing prediction')
REQUESTS = Counter('prediction_requests_total', 'Total requests')
SENTIMENT_SCORE = Histogram('sentiment_score', 'Histogram of sentiment scores', buckets=[0.0, 0.25, 0.5, 0.75, 1.0])
class TextInput(BaseModel):
text: str
class SentimentOutput(BaseModel):
text: str
sentiment: str
score: float
@app.post("/predict/naive", response_model=SentimentOutput)
async def predict_naive_sentiment(input_data: TextInput):
REQUESTS.inc()
start_time = time.time()
result = naive_classifier(input_data.text)[0]
score = result["score"]
SENTIMENT_SCORE.observe(score) # Record the sentiment score
PREDICTION_TIME.observe(time.time() - start_time)
return SentimentOutput(
text=input_data.text,
sentiment=result["label"],
score=score
)
@app.post("/predict/trained", response_model=SentimentOutput)
async def predict_trained_sentiment(input_data: TextInput):
REQUESTS.inc()
start_time = time.time()
result = trained_classifier(input_data.text)[0]
score = result["score"]
SENTIMENT_SCORE.observe(score) # Record the sentiment score
We also must update our Dockerfile to include our model files.
# Filename: Dockerfile
FROM python:3.9-slimWORKDIR /app
COPY requirements.txt .
RUN pip install -r requirements.txt
COPY main.py .
COPY ./model ./model
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]
Importantly, if you are using git, make sure that you add the pytorch_model.bin file to git lfs, so that you can push to GitHub. git lfs allows you to use version control on very large files.