retro-gap/testing.py
Michal Maciaszek cb4a1f6a4f bigrams
2020-12-08 12:10:40 +01:00

59 lines
1.9 KiB
Python

import pickle
import sys
from math import log
import regex as re
def count_prob(bigrams, unigrams):
prob = (bigrams + 1.0) / (unigrams + 1)
if prob > 1.0:
return 1.0
else:
return prob
def main():
ngrams = pickle.load(open('ngrams_2.pkl', 'rb'))
vocabulary_size = len(ngrams[1])
for line in sys.stdin:
words = re.findall(r'.*\t.*\t.* (.*?)\t(.*?) ', line.lower())[0]
#print(words)
left_word = [str(words[0])]
right_word = [str(words[1])]
probabilities = []
for word in ngrams[1].keys():
word = str(word[0])
pre_ngram = tuple(left_word + [word])
post_ngram = tuple([word] + right_word)
#print(pre_ngram)
pre_ngram_prob = count_prob(ngrams[2].get(pre_ngram, 0), ngrams[1].get(word[0],0) + vocabulary_size * 1000)
#if pre_ngram_prob>0:
post_ngram_prob = count_prob(ngrams[2].get(post_ngram, 0), ngrams[1].get(word[0],0) + vocabulary_size * 1000)
probabilities.append((word, pre_ngram_prob * post_ngram_prob))
probabilities = sorted(probabilities, key=lambda t: t[1], reverse=True)[:50]
probability = 1.0
text = ''
counter = 0
has_log_prob0 = False
for p in probabilities:
word = p[0]
prob = p[1]
if counter == 0 and (probability - prob <= 0.0):
text = word + ':' + str(log(0.95)) + ' :' + str(log(0.05))
has_log_prob0 = True
break
if counter > 0 and (probability - prob <= 0.0):
text += ':' + str(log(probability))
has_log_prob0 = True
break
text += word + ':' + str(log(prob)) + ' '
probability -= prob
counter += 1
if not has_log_prob0:
text += ':' + str(log(0.0001))
print(text)
if __name__ == '__main__':
main()