基于spark的plsa实现

PLSA.py

# coding:utf8
from pyspark import SparkContext
from pyspark import RDD
import numpy as np
from numpy.random import RandomState

import sys
reload(sys)
#设置默认编码为utf8，从spark rdd中取出中文词汇时需要编码为中文编码，否则不能保存成功
sys.setdefaultencoding('utf8')


"""
总结：
broadcast变量和需要用到broadcast变量的方法需要位于同一作用域

broadcast变量的unpersist会将存储broadcast变量的文件立即删除，
    而此时rdd并未被触发执行，当rdd执行时会发现没有broadcast变量，所以会报错，
    建议只在程序运行完成后，将broadcast变量 unpersist
"""

class PLSA:

    def __init__(self, data, sc, k, is_test=False, max_itr=1000, eta=1e-6):

        """
        init the algorithm

        :type data RDD
        :param data: 输入文章rdd，每条记录为一系列用空格分隔的词，如"我 爱 蓝天 我 爱 白云"
        :type max_itr int
        :param max_itr: 最大EM迭代次数
        :type is_test bool
        :param is_test: 是否为测试,是则rd = RandomState(1)，否则 rd = RandomState()
        :type sc SparkContext
        :param sc: spark context
        :type k int
        :param k : 主题个数
        :type eta float
        :param : 阈值，当log likelyhood的变化小于eta时，停止迭代
        :return : PLSA object
        """
        self.max_itr = max_itr
        self.k = sc.broadcast(k)
        self.ori_data = data.map(lambda x: x.split(' '))
        self.sc = sc
        self.eta = eta

        self.rd = sc.broadcast(RandomState(1) if is_test else RandomState())

    def train(self):

        #获取词汇字典 ，如{"我":1}
        self.word_dict_b = self._init_dict_()
        #将文本中词汇，转成词典中的index
        self._convert_docs_to_word_index()
        #初始化，每个主题下的单词分布
        self._init_probility_word_topic_()

        pre_l= self._log_likelyhood_()

        print "L(%d)=%.5f" %(0,pre_l)

        for i in range(self.max_itr):
            #更新每个单词主题的后验分布
            self._E_step_()
            #最大化下界
            self._M_step_()
            now_l = self._log_likelyhood_()

            improve = np.abs((pre_l-now_l)/pre_l)
            pre_l = now_l

            print "L(%d)=%.5f with %.6f%% improvement" %(i+1,now_l,improve*100)
            if improve <self.eta:
                break

    def _M_step_(self):
        """
        更新参数 p(z=k|d),p(w|z=k)
        :return: None
        """
        k = self.k
        v = self.v

        def update_probility_of_doc_topic(doc):
            """
            更新文章的主题分布
            """
            doc['topic'] = doc['topic'] - doc['topic']

            topic_doc = doc['topic']
            words = doc['words']
            for (word_index,word) in words.items():
                topic_doc += word['count']*word['topic_word']
            topic_doc /= np.sum(topic_doc)

            return {'words':words,'topic':topic_doc}

        self.data = self.data.map(update_probility_of_doc_topic)
        """
        rdd相当于一系列操作过程的结合，且前面的操作过程嵌套在后面的操作过程里，当这个嵌套超过大约60，spark会报错，
        这里每次M step都通过cache将前面的操作执行掉
        """
        self.data.cache()

        def update_probility_word_given_topic(doc):
            """
            更新每个主题下的单词分布
            """
            probility_word_given_topic = np.matrix(np.zeros((k.value,v.value)))

            words = doc['words']
            for (word_index,word) in words.items():
                probility_word_given_topic[:,word_index] += np.matrix(word['count']*word['topic_word']).T

            return probility_word_given_topic

        probility_word_given_topic = self.data.map(update_probility_word_given_topic).sum()
        probility_word_given_topic_row_sum = np.matrix(np.sum(probility_word_given_topic,axis=1))

        #使每个主题下单词概率和为1
        probility_word_given_topic = np.divide(probility_word_given_topic,probility_word_given_topic_row_sum)

        self.probility_word_given_topic = self.sc.broadcast(probility_word_given_topic)

    def _E_step_(self):
        """
        更新隐变量 p(z|w,d)-给定文章，和单词后，该单词的主题分布
        :return: None
        """
        probility_word_given_topic = self.probility_word_given_topic
        k = self.k

        def update_probility_of_word_topic_given_word(doc):
            topic_doc = doc['topic']
            words = doc['words']

            for (word_index,word) in words.items():
                topic_word = word['topic_word']
                for i in range(k.value):
                    topic_word[i] = probility_word_given_topic.value[i,word_index]*topic_doc[i]
                #使该单词各主题分布概率和为1
                topic_word /= np.sum(topic_word)
            return {'words':words,'topic':topic_doc}

        self.data = self.data.map(update_probility_of_word_topic_given_word)

    def  _init_probility_word_topic_(self):
        """
        init p(w|z=k)
        :return: None
        """
        #dict length(words in dict)
        m = self.v.value

        probility_word_given_topic = self.rd.value.uniform(0,1,(self.k.value,m))
        probility_word_given_topic_row_sum = np.matrix(np.sum(probility_word_given_topic,axis=1)).T

        #使每个主题下单词概率和为1
        probility_word_given_topic = np.divide(probility_word_given_topic,probility_word_given_topic_row_sum)

        self.probility_word_given_topic = self.sc.broadcast(probility_word_given_topic)

    def _convert_docs_to_word_index(self):

        word_dict_b = self.word_dict_b
        k = self.k
        rd = self.rd
        '''
        I wonder is there a better way to execute function with broadcast varible
        '''
        def _word_count_doc_(doc):
            wordcount ={}
            word_dict = word_dict_b.value
            for word in doc:
                if wordcount.has_key(word_dict[word]):
                    wordcount[word_dict[word]]['count'] += 1
                else:
                    #first one is the number of word occurance
                    #second one is p(z=k|w,d)
                    wordcount[word_dict[word]] = {'count':1,'topic_word': rd.value.uniform(0,1,k.value)}

            topics = rd.value.uniform(0, 1, k.value)
            topics = topics/np.sum(topics)
            return {'words':wordcount,'topic':topics}

        self.data = self.ori_data.map(_word_count_doc_)

    def _init_dict_(self):
        """
        init word dict of the documents,
        and broadcast it
        :return: None
        """
        words = self.ori_data.flatMap(lambda d: d).distinct().collect()
        word_dict = {w: i for w, i in zip(words, range(len(words)))}
        self.v = self.sc.broadcast(len(word_dict))
        return self.sc.broadcast(word_dict)

    def _log_likelyhood_(self):
        probility_word_given_topic = self.probility_word_given_topic
        k = self.k

        def likelyhood(doc):
            l = 0.0
            topic_doc = doc['topic']
            words = doc['words']

            for (word_index,word) in words.items():
                l += word['count']*np.log(np.matrix(topic_doc)*probility_word_given_topic.value[:,word_index])
            return l
        return self.data.map(likelyhood).sum()

    def save(self,f_word_given_topic,f_doc_topic):
        """
        保存模型结果 TODO 添加分布式保存结果
        :param f_word_given_topic: 文件路径，用于给定主题下词汇分布
        :param f_doc_topic: 文件路径，用于保存文档的主题分布
        :return:
        """
        doc_topic = self.data.map(lambda x:' '.join([str(q) for q in x['topic'].tolist()])).collect()
        probility_word_given_topic = self.probility_word_given_topic.value

        word_dict = self.word_dict_b.value
        word_given_topic = []

        for w,i in word_dict.items():
            word_given_topic.append('%s %s' %(w,' '.join([str(q[0]) for q in probility_word_given_topic[:,i].tolist()])))

        f1 = open (f_word_given_topic, 'w')

        for line in word_given_topic:
            f1.write(line)
            f1.write('
')
        f1.close()

        f2 = open (f_doc_topic, 'w')

        for line in doc_topic:
            f2.write(line)
            f2.write('
')
        f2.close()

调用

from PLSA import PLSA
from pyspark import SparkContext

if __name__=="__main__":
    sc = SparkContext('local')
    data = sc.textFile("E:/github/FGYML4/data/news_seg/news_seg.txt")
    plsa = PLSA(data,sc,3,max_itr=1)
    plsa.train()
    plsa.save('D:/topic_word','D:/doc_topic')