Basic Inverted Index

/*
tokenizer.go
*/

package main

import (
    "strings"
    "unicode"
)

type TokenType int

const (
    word TokenType = iota
    punct
    whitespace
)

type Token struct {
    Value string
    Type  TokenType
}

func isSeparator(r rune) bool {
    return unicode.IsSpace(r) || unicode.IsPunct(r)
}

func tokenize(input string) []Token {
    var token []Token

    var currTokenType TokenType
    var currTokenVal strings.Builder

    for _, r := range input {
        if isSeparator(r) {
            if currTokenType == word {
                tokens = append(tokens, Token{
                    Value: currTokenVal.String(),
                    Type:  currTokenType,
                })
                currTokenVal.Reset()
            }
            currTokenType = whitespace
            if unicode.IsPunct(r) {
                currTokenType = punct
            }
            tokens = append(tokens, Token{
                Value: string(r),
                Type:  currTokenType,
            })
        } else {
            if currTokenType != word {
                currTokenType = word
            }
            currTokenVal.WriteRune(r)
        }
    }

    if currTokenType == word {
        tokens = append(tokens, Token{
            Value: currTokenVal.String(),
            Type:  currTokenType,
        })
    }

    return tokens
}

/*
stemmer.go
*/

package main

import (
    "strings"

    golem "github.com/aaaton/golem/v4"
    "github.com/aaaton/golem/v4/dicts/en"
)

// Initialize the stopwords with the english stopwords, you can ask the LLM to generate
// this for you. Map them to empty struct
// because empty struct are of size 0.
var stopwords = map[string]struct{}{
    // initialize here.
}

func removeStopWords(tokens []Token) []Token {
    result := []Token{}
    for _, token := range tokens {
        _, ok := stopwords.[token.Value]
        if !ok {
            result = append(result, token)
        }
    }
    return result
}

func removePunctAndSpace(tokens []Token) []Token {
    result := []Token{}
    for _, token := range tokens {
        if token.Type == word {
            result = append(result, token)
        }
    }
    return result
}

func lemmetizedTokens(tokens []Token) []Token {
    lemmer, err := golem.New(en.New())
    if err != nil {
        panic(err)
    }
    for i, token := range tokens {
        token.Value = lemmer.Lemma(token.Value)
        tokens[i] = token
    }
    return tokens
}

func stemmer(tokens []Token) []Token {
    for i, token := range tokens {
        tokens[i] = Token{
            Value: strings.ToLower(token.Value),
            Type:  token.Type,
        }
    }
    tokens = removePunctAndSpace(tokens)
    tokens = removeStopWords(tokens)

    return lemmetizedTokens(tokens)
}

/*
indexer.go
*/

package main

import (
    "encoding/gob"
    "errors"
    "os"
    "path/filepath"
    "slices"
)

type IR struct {
    Index map[string][]int
}

func NewIR() *IR {
    return &IR{
        Index: make(map[string][]int),
    }
}

// save the data into a file, for persistant storage.
func (i *IR) serializer(indexPath string) {
    f, err := os.OpenFile(indexPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
    if err != nil {
        panic(err)
    }

    encoder := gob.NewEncoder(f)

    err = encoder.Encode(i)
    if err != nil {
        panic(err)
    }
}

// loads any data that might be already indexed.
func (i *IR) deserializer(indexFile string) {
    f, err := os.Open(indexFile)
    if err != nil {
        if errors.Is(err, os.ErrNotExist) {
            return
        }
        panic(err)
    }

    decoder := gob.NewDecoder(f)

    err = decoder.Decode(i)
    if err != nil {
        panic(err)
    }
}

func indexer(tokens []Token, docID int) *IR {
    cwd, err := os.Getwd()
    if err != nil {
        panic(err)
    }
    indexFile := filepath.Join(cwd, "invert.index")

    ir := NewIR()

    ir.deserializer(indexFile)

    for _, token := range tokens {
        list, found := ir.Index[token.Value]
        // if the word is not present in the index, then we add
        // it to the index, with docID it was found in and
        // go to next iteration
        if !found {
            ir.Index[token.Value] = []int{docID}
            continue
        }

        // check if the docID is present, if it not then add
        // the docID in the slice of docID's
        present := slices.Contains(list, docID)
        if !present {
            ir.Index[token.Value] = append(list, docID)
        }
    }

    ir.serializer(indexFile)

    return ir
}

/*
main.go
*/

package main

import (
	"flag"
    "fmt"
    "sort"
	"os"
)

func main() {
	flag.String("init", "", "Initialize the data and create an Inverted index.")

	flag.Parse()

	if len(os.Args) < 2 {
		flag.Usage()
		return
	}

	switch os.Args[1] {
	case "init":
		initializeData()
}


func initializeData() {
	inputs := []struct {
		ID   int
		Data string
	}{
		{
			ID:   1,
			Data: "I did enact Julius Caesar: I was killed i' the Capitol; Brutus killed me",
		},
		{
			ID:   2,
			Data: "So let it be with Caesar. The noble Brutus hath told you Caesar was ambitious",
		},
	}

	for _, input := range inputs {
		tokens := tokenize(input.Data)

		fmt.Println("Original input: ", input)
		fmt.Println("Tokenized output:")

		for i, token := range tokens {
			fmt.Printf("%d: %s (%d)\n", i+1, token.Value, token.Type)
		}

		stemmedTokens := stemmer(tokens)
		fmt.Println("Stemmed output:")

		for i, token := range stemmedTokens {
			fmt.Printf("%d: %s (%d)\n", i+1, token.Value, token.Type)
		}

		fmt.Println("Starting indexing the input.")
		ir := indexer(stemmedTokens, input.ID)
		fmt.Println("Finished indexing the input.")

		printSortedMap(ir)
	}
}

func printSortedMap(ir *IR) {
	fmt.Println("Printing the index.")

	keys := make([]string, 0, len(ir.Index))

	for key := range ir.Index {
		keys = append(keys, key)
	}
	sort.Strings(keys)

	for _, k := range keys {
		fmt.Printf("word: %s, list: %v\n", k, ir.Index[k])
	}
}