Create tf-idf.py

tf-idf.py

@@ -0,0 +1,26 @@
+#!/usr/bin/python
+# -*- coding: latin-1 -*-
+
+import math
+from operator import itemgetter
+
+def tfidf(word, doc, docList):
+    tf = doc.split(None).count(word) / float(len(doc.split(None)))
+
+    docs_in_word = 0
+
+    for doc in docList:
+        if doc.split(None).count(word) > 0:
+            docs_in_word +=1
+
+    idf = math.log(len(docList)/ docs_in_word)
+
+    return tf * idf
+
+def top_keywords(doc, docList, n=5):
+    d = {}
+    for word in set(doc.split(None)):
+        d[word] = tfidf(word, doc, docList)
+
+    sorted_d = sorted(d.items(), key=itemgetter(1), reverse=True)
+    return [w[0] for w in sorted_d[:n]]