Writing a simple text classifier in Python

Many people want to write a simple text classifier, but get lost in tons of machine learning books and complex mathematical formulas. Today I will show you a relatively simple Python classification example that works in a simple and straightforward way.

In general, the scheme of machine learning is as follows:

  1. Create a program using special libraries

  2. Teach the program big data set, that is, they give it a data set correctly labeled by a person so that the program learns and starts working on other similar data sets

  3. Save and use the trained model

In my example, I combined the training of the program along with its testing, so every time it is launched, it first learns (which takes some time, do not think that the program is frozen).

In addition, in my example, I will use a very small data set; in real projects, the data set should be at least several thousand rows.

So, we will create a program that tries to guess what type of content to find for the user, in response to his question. For example, if the user asks “Who is Viktor Tsoi”, the program should answer “Personal Information”. It won’t answer the question, just classifies the question into a specific category. The training data set is very small, so we will test on more or less similar questions.

Create a training dataset

First, let’s create a training dataset. I have a plain text file that contains lines separated by the @ sign. Name the file model.txt and drop it next to the script. Use UTF-8 encoding.

Что такое патока @ Информация о предмете
Где живут медведи @ Местоположение
Почему небо голубое @ Информация о предмете
Как сварить борщ @ Инструкция
Как научиться летать @ Инструкция
Как рано просыпаться @ Инструкция
Как обрести счастье @ Инструкция
Как убить таракана @ Инструкция
Сколько звезд на небе @ Количество
Какая высота у Эйфелевой башни @ Количество
Почему идет дождь @ Информация о предмете
В каких фильмах играл Ди Каприо @ Информация о предмете
Когда родился Эйнштейн @ Дата рождения
Когда началась Первая мировая война @ Дата
В каком году был построен Титаник @ Дата
Когда начал править Иван IV @ Дата
Когда была битва под Аустерлицем @ Дата
В каких годах была «Семилетняя война» @ Дата
Когда начал править Наполеон Бонапарт @ Дата
Когда закончилось правление Наполеона Бонапарта @ Дата
В каких годах правил князь Святослав @ Дата
В каком году умер Ярослав Мудрый @ Дата
В каком году появляется первое упоминание о Москве @ Дата
В каком году был основан Санкт-Петербург @ Дата
Когда было Ледовое побоище @ Дата
Когда была Куликовская битва @ Дата
В каком году начали печатать книги на Руси @ Дата
Когда была Полтавская битва @ Дата
В каком году закончилась Первая мировая война @ Дата
В каком году был заключен Тильзитский мир @ Дата
В каком году были написаны «Отцы и дети» @ Дата
Когда отмечается День Учителя @ Дата
Когда отмечается День Матери @ Дата
Когда отменили крепостное право @ Дата
Что произошло 12 июня 1812 @ Дата
Что произошло 20 ноября 1805 @ Дата
Что произошло 6 июня 1799 @ Дата
Что произошло 1 апреля 1809 @ Дата
Какой город начинает упоминатся с 1147 грода @ Дата
Кто такой Эйнштейн @ Информация о личности
Что такое «дождь» @ Информация о предмете
Кто основал Санкт-Петербург @ Имя
Кто такой Христофор Колумб @ Информация о личности
Как звали Колумба @ Имя
Кто написал произведение «Преступление и наказание» @ Имя
Кто написал стихотворение «Тройка» @ Имя
Какова средняя глубина Тихого океана @ Количество
Какова высота Биг Бена @ Количество
Где находится Сиднейская Опера @ Местоположение
Где находится самое большое озеро @ Местоположение
В какой стране столицей является город Москва @ Местоположение

Installing the required libraries

We need the PyStemmer library. You can install it like this

pip3 install pystemmer

We write a program

Now I will show the script itself for classifying questions.

import sys
import numpy as np
import pickle
import re
from Stemmer import Stemmer
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import SGDClassifier
from sklearn.pipeline import Pipeline
from sklearn.metrics import accuracy_score

# очистка текста с помощью regexp приведение слов в инфинитив и нижний регистр, замена цифр

def text_cleaner(text):
    text = text.lower() # приведение в lowercase 
    stemmer = Stemmer('russian')
    text=" ".join( stemmer.stemWords( text.split() ) ) 
    text = re.sub( r'\b\d+\b', ' digit ', text ) # замена цифр 
    return  text 

# загрузка данных из файла model.txt

def load_data():   
    data = {'text':[],'tag':[]}
    for line in open('model.txt'):
        if not('#' in line):
            row = line.split("@") 
            data['text'] += [row[0]]
            data['tag'] += [row[1]]
    return data

# Обучение 

def train_test_split(data, validation_split = 0.1):
    sz = len(data['text'])
    indices = np.arange(sz)

    X = [data['text'][i] for i in indices]
    Y = [data['tag'][i] for i in indices]
    nb_validation_samples = int( validation_split * sz )

    return { 
        'train': {'x': X[:-nb_validation_samples], 'y': Y[:-nb_validation_samples]},
        'test': {'x': X[-nb_validation_samples:], 'y': Y[-nb_validation_samples:]}

# - - - -

def openai():
    data = load_data()
    D = train_test_split(data)
    text_clf = Pipeline([
                    ('tfidf', TfidfVectorizer()),
                    ('clf', SGDClassifier(loss="hinge")),
    text_clf.fit(D['train']['x'], D['train']['y'])
    predicted = text_clf.predict( D['train']['x'] )
# Начало тестирования программы
    z = input("Введите вопрос без знака вопроса на конце: ")
    zz = []
    predicted = text_clf.predict(zz) 
# - - - -
if __name__ == '__main__':

Let’s run our program, wait for it to learn (a couple of minutes) and prompt us to enter a question. Let’s introduce a question that is not in the training file – “Who Invented the Rocket“. You need to write without a question mark at the end. In response, the program will issue “Name“, so she determined that our question implies an answer in which there should be someone’s name. Note that this question was not in the model.txt file, but the program accurately determined what we mean, based on the training sample.

This classifier is very simple. You can enter absolutely any data that needs to be classified into the model.txt file, it does not have to be questions and their categories, as in my example.

Similar Posts

Leave a Reply

Your email address will not be published. Required fields are marked *