assignment / app /classifier.py
sameer9's picture
Add FastAPI app and Docker configuration
e5b2387
from sklearn.linear_model import LogisticRegression
from sklearn.feature_extraction.text import TfidfVectorizer
def train_classifier(dummy_data):
vectorizer = TfidfVectorizer()
train_texts, train_labels = zip(*dummy_data)
train_vectors = vectorizer.fit_transform(train_texts)
classifier = LogisticRegression()
classifier.fit(train_vectors, train_labels)
return classifier, vectorizer
def classify_text(text: str, classifier, vectorizer) -> str:
try:
transformed_data = vectorizer.transform([text])
category = classifier.predict(transformed_data)[0]
return category
except Exception as e:
return str(e)