Utilizando un modelo personalizado
Si bien el proceso de construcción y ajuste de un modelo no es la intención de este proyecto, es importante comprender cómo se puede agregar un modelo a este proceso.
# 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()
Para asegurarnos de que podemos consultar nuestro nuevo modelo que hemos entrenado, debemos actualizar algunos de nuestros archivos existentes. Por ejemplo, en main.py ahora usamos el modelo de ./model y cárguelo como un modelo previamente entrenado. Además, a modo de comparación, agregamos que ahora tenemos dos puntos finales para usar, /predict/naive y 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
También debemos actualizar nuestro Dockerfile para incluir nuestros archivos de modelo.
# 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"]
Es importante destacar que si estás usando git, asegúrate de agregar el pytorch_model.bin archivo a git lfspara que puedas enviar a GitHub. git lfs te permite usar el control de versiones en archivos muy grandes.