Search code examples
pythonpandasscikit-learnnaivebayesmultinomial

How to add a predicted-data column to my dataframe?


I'm using naive bayes to predict country name from a set of addresses, I tried this

import re
import numpy as np
import pandas as pd
from sklearn.naive_bayes import MultinomialNB
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.preprocessing import LabelEncoder
from sklearn.metrics import accuracy_score
def normalize_text(s):
    s = s.lower()
    s = re.sub('\s\W',' ',s)
    s = re.sub('\W\s',' ',s)
    s = re.sub('\s+',' ',s)
    return(s)
df['TEXT'] = [normalize_text(s) for s in df['Full_Address']]

# pull the data into vectors
vectorizer = CountVectorizer()
x = vectorizer.fit_transform(df['TEXT'])

encoder = LabelEncoder()
y = encoder.fit_transform(df['CountryName'])

# split into train and test sets
x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2)

nb = MultinomialNB()
nb.fit(x_train, y_train)
y_predicted = nb.predict(x_test)

so what I want is to add another column to my dataframe with the predicted Country name, how can I achieve that?

Update:

df['Predicted'] = nb.predict(x)

         CountryName                                       Full_Address  \
8913       Indonesia  EJIP Industrial Park Plot 1E-2, Sukaresmi, Cik...   
7870   United States    360 Thelma Street, Sandusky, Michigan 48471 USA   
32037          China  1027, 26/F, Zhao Feng Mansion, Chang Ning Road...   
38769    New Zealand  NZ - 164 ST. ASAPH STREET, \tCHRISTCHURCH 8011...   
46639          India  301-306, Sahajanand Trade Center, Opp. Kothawa...   

                                                    TEXT  Predicted  
8913   ejip industrial park plot 1e-2 sukaresmi cikar...         66  
7870       360 thelma street sandusky michigan 48471 usa        169  
32037  1027 26/f zhao feng mansion chang ning road sh...         30  
38769  nz 164 st asaph street christchurch 8011 new z...        112  
46639  301-306 sahajanand trade center opp kothawala ...         65

Solution

  • You should use the inverse of encoder.fit_transform on the predicted values of y, applied to the output of the model. So something like

    df['Predicted'] = encoder.inverse_transform(nb.predict(x))
    

    This assumes that nb.predict(x)'s output is a list of integers (rather than a list of lists) -- you may have do so some reshaping if it is not. Since I cannot run your code without access to df I can't really say