category_test.py

# coding=utf-8

import sys
reload(sys)
sys.setdefaultencoding('utf-8')

import re
import csv
import requests
import jieba
import jieba.analyse
from bs4 import BeautifulSoup as bs
from tgrocery import Grocery
from train_txt_5 import train_src


def artical_content(url):
 rsp = requests.get(url)
 rsp.encoding = 'utf-8'
 html = bs(rsp.text, 'lxml')
 # artical = html.select('#sina_keyword_ad_area2')
 artical = html.select('.h1_tit')
 if len(artical) > 0:
 content = artical[0].text
 else:
 content = html.select('.SG_txta')[0].text
 # content = ''
 return content.strip()


def category(urls):
 for url in urls:
 artical = artical_content(url)
 cate2 = new_grocery.predict(artical)
 yield cate2.predicted_y


def unzip(seq, L=None):
 if L is None:
 L = []
 for s in seq:
 if not isinstance(s, (list, )):
 L.append(s)
 else:
 L.extend(unzip(s))
 return L


if __name__ == "__main__":

 grocery = Grocery('sample')
 grocery.train(train_src)
 grocery.save()
 new_grocery = Grocery('sample')
 new_grocery.load()

 L1 = []
 with open('/home/mouse/Downloads/female.csv', 'r') as f1:
 f1_csv = csv.reader(f1)
 for row in f1_csv:
 L1.append(row[0])
 # print(len(L1))

 cate = category(L1)
 i = 1
 with open('/home/mouse/infoss.csv', 'w') as f:
 f_csv = csv.writer(f)
 for row in zip(L1, cate):
 f_csv.writerow(unzip(row))
 print 'Writing now, please waiting...{}'.format(str(i))
 i += 1
 print 'Done'