[Python] TextRank 구현 코드

앞선글(TextRank 기법을 이용한 핵심 어구 추출 및 텍스트 요약) 에서 텍스트 랭크 알고리즘을 소개하면서 직접 실험해보고자 파이썬 코드를 짰었는데요, 정리해서 공유해드립니다. 페이지랭크 계산에 networkx 패키지를 사용하므로, networkx가 설치되어있는지 확인부탁드립니다.

* 의존 패키지: networkx

import networkx import re class RawSentence: def __init__(self, textIter): if type(textIter) == str: self.textIter = textIter.split('\n') else: self.textIter = textIter self.rgxSplitter = re.compile('([.!?:](?:["\']|(?![0-9])))') def __iter__(self): for line in self.textIter: ch = self.rgxSplitter.split(line) for s in map(lambda a, b: a + b, ch[::2], ch[1::2]): if not s: continue yield s class RawSentenceReader: def __init__(self, filepath): self.filepath = filepath self.rgxSplitter = re.compile('([.!?:](?:["\']|(?![0-9])))') def __iter__(self): for line in open(self.filepath, encoding='utf-8'): ch = self.rgxSplitter.split(line) for s in map(lambda a, b: a + b, ch[::2], ch[1::2]): if not s: continue yield s class RawTagger: def __init__(self, textIter, tagger = None): if tagger: self.tagger = tagger else : from konlpy.tag import Komoran self.tagger = Komoran() if type(textIter) == str: self.textIter = textIter.split('\n') else: self.textIter = textIter self.rgxSplitter = re.compile('([.!?:](?:["\']|(?![0-9])))') def __iter__(self): for line in self.textIter: ch = self.rgxSplitter.split(line) for s in map(lambda a,b:a+b, ch[::2], ch[1::2]): if not s: continue yield self.tagger.pos(s) class RawTaggerReader: def __init__(self, filepath, tagger = None): if tagger: self.tagger = tagger else : from konlpy.tag import Komoran self.tagger = Komoran() self.filepath = filepath self.rgxSplitter = re.compile('([.!?:](?:["\']|(?![0-9])))') def __iter__(self): for line in open(self.filepath, encoding='utf-8'): ch = self.rgxSplitter.split(line) for s in map(lambda a,b:a+b, ch[::2], ch[1::2]): if not s: continue yield self.tagger.pos(s) class TextRank: def __init__(self, **kargs): self.graph = None self.window = kargs.get('window', 5) self.coef = kargs.get('coef', 1.0) self.threshold = kargs.get('threshold', 0.005) self.dictCount = {} self.dictBiCount = {} self.dictNear = {} self.nTotal = 0 def load(self, sentenceIter, wordFilter = None): def insertPair(a, b): if a > b: a, b = b, a elif a == b: return self.dictBiCount[a, b] = self.dictBiCount.get((a, b), 0) + 1 def insertNearPair(a, b): self.dictNear[a, b] = self.dictNear.get((a, b), 0) + 1 for sent in sentenceIter: for i, word in enumerate(sent): if wordFilter and not wordFilter(word): continue self.dictCount[word] = self.dictCount.get(word, 0) + 1 self.nTotal += 1 if i - 1 >= 0 and (not wordFilter or wordFilter(sent[i-1])): insertNearPair(sent[i-1], word) if i + 1 < len(sent) and (not wordFilter or wordFilter(sent[i+1])): insertNearPair(word, sent[i+1]) for j in range(i+1, min(i+self.window+1, len(sent))): if wordFilter and not wordFilter(sent[j]): continue if sent[j] != word: insertPair(word, sent[j]) def loadSents(self, sentenceIter, tokenizer = None): import math def similarity(a, b): n = len(a.intersection(b)) return n / float(len(a) + len(b) - n) / (math.log(len(a)+1) * math.log(len(b)+1)) if not tokenizer: rgxSplitter = re.compile('[\\s.,:;-?!()"\']+') sentSet = [] for sent in filter(None, sentenceIter): if type(sent) == str: if tokenizer: s = set(filter(None, tokenizer(sent))) else: s = set(filter(None, rgxSplitter.split(sent))) else: s = set(sent) if len(s) < 2: continue self.dictCount[len(self.dictCount)] = sent sentSet.append(s) for i in range(len(self.dictCount)): for j in range(i+1, len(self.dictCount)): s = similarity(sentSet[i], sentSet[j]) if s < self.threshold: continue self.dictBiCount[i, j] = s def getPMI(self, a, b): import math co = self.dictNear.get((a, b), 0) if not co: return None return math.log(float(co) * self.nTotal / self.dictCount[a] / self.dictCount[b]) def getI(self, a): import math if a not in self.dictCount: return None return math.log(self.nTotal / self.dictCount[a]) def build(self): self.graph = networkx.Graph() self.graph.add_nodes_from(self.dictCount.keys()) for (a, b), n in self.dictBiCount.items(): self.graph.add_edge(a, b, weight=n*self.coef + (1-self.coef)) def rank(self): return networkx.pagerank(self.graph, weight='weight') def extract(self, ratio = 0.1): ranks = self.rank() cand = sorted(ranks, key=ranks.get, reverse=True)[:int(len(ranks) * ratio)] pairness = {} startOf = {} tuples = {} for k in cand: tuples[(k,)] = self.getI(k) * ranks[k] for l in cand: if k == l: continue pmi = self.getPMI(k, l) if pmi: pairness[k, l] = pmi for (k, l) in sorted(pairness, key=pairness.get, reverse=True): print(k[0], l[0], pairness[k, l]) if k not in startOf: startOf[k] = (k, l) for (k, l), v in pairness.items(): pmis = v rs = ranks[k] * ranks[l] path = (k, l) tuples[path] = pmis / (len(path) - 1) * rs ** (1 / len(path)) * len(path) last = l while last in startOf and len(path) < 7: if last in path: break pmis += pairness[startOf[last]] last = startOf[last][1] rs *= ranks[last] path += (last,) tuples[path] = pmis / (len(path) - 1) * rs ** (1 / len(path)) * len(path) used = set() both = {} for k in sorted(tuples, key=tuples.get, reverse=True): if used.intersection(set(k)): continue both[k] = tuples[k] for w in k: used.add(w) #for k in cand: # if k not in used or True: both[k] = ranks[k] * self.getI(k) return both def summarize(self, ratio = 0.333): r = self.rank() ks = sorted(r, key=r.get, reverse=True)[:int(len(r)*ratio)] return ' '.join(map(lambda k:self.dictCount[k], sorted(ks)))

사용법은 어렵지 않습니다. 다음과 같이 쓰면 됩니다.

TextRank 생성자는 인수로 window, coef, threshold를 받습니다.

* window : 문맥으로 사용할 단어의 개수. 기본값 5로 주면 특정 단어의 좌우 5개씩, 총 10개 단어를 문맥으로 사용합니다.

* coef : 동시출현 빈도를 weight에 반영하는 비율입니다. 기본값은 1.0로, 동시출현 빈도를 weight에 전부 반영합니다. 0.0일 경우 빈도를 반영하지 않고 모든 간선의 weight을 1로 동일하게 간주합니다.

* threshold: 문서 요약시 관련있는 문장으로 여길 최소 유사도값. 기본값은 0.005이고, 이 값보다 작은 유사도를 가지는 문장쌍은 관련없는문장으로 처리합니다.

문서요약시에는 loadSents로 문장을 읽어들이고, summarize로 결과를 얻습니다.

tr = TextRank() print('Load...') from konlpy.tag import Komoran tagger = Komoran() stopword = set([('있', 'VV'), ('하', 'VV'), ('되', 'VV') ]) tr.loadSents(RawSentenceReader('test2.txt'), lambda sent: filter(lambda x:x not in stopword and x[1] in ('NNG', 'NNP', 'VV', 'VA'), tagger.pos(sent))) print('Build...') tr.build() ranks = tr.rank() for k in sorted(ranks, key=ranks.get, reverse=True)[:100]: print("\t".join([str(k), str(ranks[k]), str(tr.dictCount[k])])) print(tr.summarize(0.2))

키워드 추출시에는 load로 문장을 읽어들이고, extract로 키워드를 추출합니다.

tr = TextRank(window=5, coef=1) print('Load...') stopword = set([('있', 'VV'), ('하', 'VV'), ('되', 'VV'), ('없', 'VV') ]) tr.load(RawTaggerReader('test2.txt'), lambda w: w not in stopword and (w[1] in ('NNG', 'NNP', 'VV', 'VA'))) print('Build...') tr.build() kw = tr.extract(0.1) for k in sorted(kw, key=kw.get, reverse=True): print("%s\t%g" % (k, kw[k]))