import sys
import datetime
from os.path import isdir, isfile
from corpus import Corpus
from gensim import models


class W2VCorpus(Corpus):
    # TODO: the clipped corpus class doesn't implement get_texts(), investigate this
    def __init__(self, directory=None, dictionary=None, corpus=None, w2v_model=None, max_docs=None, **kwargs):
        Corpus.__init__(self, directory=directory, dictionary=dictionary, corpus=corpus, max_docs=max_docs)
        self.dict_loc = dictionary
        self.vec_loc = corpus
        self.model = None
        start_time = datetime.datetime.now()
        if not w2v_model:
            # Todo: Tweak the default parameter
            self.model = models.Word2Vec(self.docs.get_texts(), size=100, window=5, min_count=5, workers=4)
        else:
            self.model = models.Word2Vec.load(w2v_model)
        end_time = datetime.datetime.now()
        self.train_time = end_time - start_time
        return

    def similarity(self, word1, word2):
        if self.model:
            return self.model.similarity(word1, word2)
        else:
            # Todo: Raise exception?
            return None

    def save(self, dictionary_file="w2v_corpus.dict", corpus_file="corpus.mm", sup_file="vector.w2vs"):
        Corpus.save(self, dictionary_file, corpus_file)
        if self.model: self.model.save(sup_file)

    @classmethod
    def load(cls, dictionary_file=None, corpus_file=None, sup_file=None):
        return cls(dictionary=dictionary_file, corpus=corpus_file, w2v_model=sup_file)

    def clip_corpus(self, max_docs=None):
        # Override the default clip_corpus function to do nothing.
        pass


def main():
    w2v = "w2vcorpus.w2v"
    if len(sys.argv) is 4 and isdir(sys.argv[1]):
        if not isfile(w2v):
            corpus = W2VCorpus(directory=sys.argv[1])
            corpus.save(dictionary_file=sys.argv[2], corpus_file=sys.argv[3] ,sup_file=w2v)
        else:
            corpus = W2VCorpus.load(dictionary_file=sys.argv[2], corpus_file=sys.argv[3], sup_file=w2v)
        time = corpus.get_train_time()
        print "W2V Train Time:\t" + str(time)
    else:
        print "Corpus requires directory as an argument."

if __name__ == "__main__": main()
@@ -79,16 +79,29 @@
 t Corpus
+, PaperCorpus
 %0Afrom ge
@@ -141,101 +141,17 @@
 pus(
-Corpus):%0A    # TODO: the clipped corpus class doesn't implement get_texts(), investigate this
+object):%0A
 %0A   
@@ -189,38 +189,8 @@
 one,
- dictionary=None, corpus=None,
  w2v
@@ -240,106 +240,84 @@
     
-Corpus.__init__(self, directory=directory, dictionary=dictionary, corpus=corpus, max_
+docs = Corpus.get_docs(directory)%0A        if max_docs: docs = 
 docs
-=
+%5B:
 max_docs
 )%0A  
@@ -312,17 +312,17 @@
 max_docs
-)
+%5D
 %0A       
@@ -332,84 +332,31 @@
 lf.d
-ict_l
 oc
+s
  = 
-dictionary%0A        self.vec_loc = corpus%0A        self.model = None
+PaperCorpus(docs)
 %0A   
@@ -499,28 +499,27 @@
 el = models.
-Word
+Doc
 2Vec(self.do
@@ -520,21 +520,16 @@
 elf.
-docs.get_text
+_gen_doc
 s(),
@@ -623,12 +623,11 @@
 els.
-Word
+Doc
 2Vec
@@ -775,28 +775,19 @@
 y(self, 
-word1, word2
+doc
 ):%0A     
@@ -838,31 +838,24 @@
 del.
+most_
 similar
-ity(word1, word2
+(doc
 )%0A  
@@ -1036,64 +1036,8 @@
 %22):%0A
-        Corpus.save(self, dictionary_file, corpus_file)%0A
     
@@ -1273,121 +1273,296 @@
 def 
-clip_corpus(self, max_docs=None):%0A        # Override the default clip_corpus function to do nothing.%0A        pass
+_gen_docs(self,):%0A        docs = self.docs.get_texts()%0A        for doc_id, doc in enumerate(docs):%0A            yield models.doc2vec.LabeledSentence(words=doc, labels=%5B'DOC_%25s' %25 doc_id%5D)%0A%0A    def run_query(self, query, index_location, best_matches):%0A        return self.similarity(query)%0A
 %0A%0A%0Ad

Switch from using word2vec to doc2vec