[Python] 3변수 상호정보량을 활용한 연어 추출 코드

최근 연어 추출과 관련하여 몇몇 문의가 있었어서 '다변수 정규화 상호정보량과 연어 추출'에서 다뤘던 다변수 상호정보량을 계산하는 파이썬 코드를 공유해드립니다.

class PMI3: def __init__(self, **kargs): self.dictCount = {} self.dictTriCount = {} self.nTotal = 0 def train(self, sentenceIter, weight = 1): for sent in sentenceIter: self.nTotal += len(sent) for word in sent: self.dictCount[word] = self.dictCount.get(word, 0) + weight for a, b, c in zip(sent[:-2], sent[1:-1], sent[2:]): self.dictTriCount[a, b, c] = self.dictTriCount.get((a, b, c), 0) + weight def getCoOccurrence(self, a, b, c): return self.dictTriCount.get((a, b, c), 0) def getPMI(self, a, b, c): import math co = self.getCoOccurrence(a, b, c) if not co: return None return math.log(float(co) * self.nTotal * self.nTotal / self.dictCount[a] / self.dictCount[b] / self.dictCount[c]) def getNPMI(self, a, b, c): import math abc = self.getPMI(a, b, c) if abc == None: return -1 return abc / (2 * math.log(self.nTotal / self.getCoOccurrence(a, b, c))) def getPMIDict(self, minNum = 5): ret = {} for a, b, c in self.dictTriCount: if self.dictTriCount[a, b, c] < minNum: continue ret[a, b, c] = self.getPMI(a, b, c) return ret def getNPMIDict(self, minNum = 5): ret = {} for a, b, c in self.dictTriCount: if self.dictTriCount[a, b, c] < minNum: continue ret[a, b, c] = self.getNPMI(a, b, c) return ret

getPMI 및 getNPMI 메서드가 연속된 3단어 a,b,c의 상호정보량 및 정규화 상호정보량을 계산해줍니다. 전체 a-b-c 셋에 대한 상호정보량 값을 알고 싶다면 getPMIDict, getNPMIDict 메서드를 사용하시면 됩니다. 최소 minNum 이상 등장한 단어 셋에 대해 각각 상호정보량 및 정규화 상호정보량을 계산하여 dict로 돌려줍니다.

이전의 상호정보량 계산 코드와 마찬가지로 다음과 같이 이터레이터를 선언하여 사용하면 편리합니다.

# 행 별로 문장이 구분되어있고, 탭으로 단어가 구분된 텍스트 파일을 읽어줍니다. class SentenceReader: def __init__(self, filepath): self.filepath = filepath def __iter__(self): for line in open(self.filepath, encoding='utf-8'): yield list(line.split('\t'))

실제 사용은 다음과 같이 하면 되겠죠?

pc = PMI3() pc.train(SentenceReader('test.txt')) # test.txt 파일을 읽어와 분석합니다. # 5회 이상 등장하는 연어에 대해 정규화 상호정보량을 계산하여 내림차순으로 출력합니다. res = pc.getNPMIDict() for a, b, c in sorted(res, key=res.get, reverse=True): print(a, b, c, res[a, b, c])