Inference for BATM about neural_topic_models HOT 2 CLOSED

def inference(args):
    device = torch.device('cpu')
    docSet = DocDataset(args.taskname, no_below=args.no_below,
                        no_above=args.no_above, rebuild=False, use_tfidf=args.use_tfidf, lang="en")

    voc_size = docSet.vocabsize

    model = BATM(bow_dim=voc_size, n_topic=args.n_topic,
                 device=device, taskname=args.taskname)

    checkpoint = torch.load("<checkpoint_path>")
    model.encoder.load_state_dict(checkpoint['encoder'])
    model.id2token = {v: k for k, v in docSet.dictionary.token2id.items()}

    bows = [docSet[i][1] for i in range(len(docSet))]
    bows = torch.stack(bows)

    results = model.encoder.e(bows)
    results = torch.argmax(results, dim=1)

    return results