Ejemplo n.º 1
0
func makeExpr(addList, subList string) (word2vec.Expr, error) {
	if addList == "" && subList == "" {
		return word2vec.Expr{}, fmt.Errorf("must specify 'add' and/or 'sub' component for each target vector; see -h for more details")
	}

	result := word2vec.Expr{}
	if addList != "" {
		for _, w := range strings.Split(addList, ",") {
			result.Add(1, w)
		}
	}
	if subList != "" {
		for _, w := range strings.Split(subList, ",") {
			result.Add(-1, w)
		}
	}
	return result, nil
}
Ejemplo n.º 2
0
func main() {
	flag.Parse()

	if path == "" {
		fmt.Println("must specify -model; see -h for more details")
		os.Exit(1)
	}

	if addList == "" && subList == "" && multiQuery == "" {
		fmt.Println("must specify -add, -sub, or -words; see -h for more details")
		os.Exit(1)
	}

	f, err := os.Open(path)
	if err != nil {
		fmt.Printf("error opening binary model data file: %v\n", err)
		os.Exit(1)
	}
	defer f.Close()

	m, err := word2vec.FromReader(f)
	if err != nil {
		fmt.Printf("error reading binary model data: %v\n", err)
		os.Exit(1)
	}

	// TODO(dhowden): Tidy this up, it's rather hacked in here!
	if multiQuery != "" {
		var exprs []word2vec.Expr

		multiWords := strings.Split(multiQuery, ",")
		for _, w := range multiWords {
			e := word2vec.Expr{}
			e.Add(1, w)

			exprs = append(exprs, e)
		}

		before := time.Now()
		res, err := word2vec.MultiCosN(m, exprs, 10)
		if err != nil {
			fmt.Printf("error retrieving multi cos: %v", err)
			os.Exit(1)
		}
		fmt.Println("Total time:", time.Since(before))
		fmt.Println(res)
		return
	}

	expr := word2vec.Expr{}
	if addList != "" {
		word2vec.Add(expr, 1, strings.Split(addList, ","))
	}
	if subList != "" {
		word2vec.Add(expr, -1, strings.Split(subList, ","))
	}

	if verbose {
		fmt.Printf("Expr: %#v\n", expr)
	}

	if verbose {
		v, err := expr.Eval(m)
		if err != nil {
			fmt.Printf("error creating target vector: %v", err)
			os.Exit(1)
		}
		fmt.Printf("Target vector: %#v\n", v)
	}

	before := time.Now()
	pairs, err := m.CosN(expr, n)
	if err != nil {
		fmt.Printf("error finding most similar: %v\n", err)
		os.Exit(1)
	}
	after := time.Now()

	if verbose {
		fmt.Println("Total time: ", after.Sub(before))
	}

	for _, k := range pairs {
		fmt.Printf("%9f\t%#v\n", k.Score, k.Word)
	}
}