github.com/goplus/llgo@v0.8.3/_demo/llama2-c/run.go (about)

     1  package main
     2  
     3  import (
     4  	"github.com/goplus/llgo/c"
     5  	"github.com/goplus/llgo/c/llama2"
     6  )
     7  
     8  func main() {
     9  	var prompt *c.Char = c.Str("Once upon a time")
    10  	var checkpointPath *c.Char = c.Str("stories15M.bin")
    11  	var tokenizerPath *c.Char = c.Str("tokenizer.bin")
    12  	var temperature, topp c.Float = 1.0, 0.9
    13  	var steps c.Int = 256
    14  	var rngSeed uint64 = uint64(c.Time(nil))
    15  
    16  loop: // parse command line arguments
    17  	for {
    18  		switch c.Getopt(c.Argc, c.Argv, c.Str("m:")) {
    19  		case 'm':
    20  			checkpointPath = c.Optarg
    21  			c.Fprintf(c.Stderr, c.Str("==> use model: %s\n"), checkpointPath)
    22  		case -1:
    23  			break loop
    24  		}
    25  	}
    26  	if c.Optind < c.Argc {
    27  		prompt = c.Index(c.Argv, c.Optind)
    28  		c.Fprintf(c.Stderr, c.Str("==> prompt: %s\n"), prompt)
    29  	}
    30  
    31  	// build the Transformer via the model .bin file
    32  	var transformer llama2.Transformer
    33  	llama2.BuildTransformer(&transformer, checkpointPath)
    34  
    35  	// build the Tokenizer via the tokenizer .bin file
    36  	var tokenizer llama2.Tokenizer
    37  	llama2.BuildTokenizer(&tokenizer, tokenizerPath, transformer.Config.VocabSize)
    38  
    39  	// build the Sampler
    40  	var sampler llama2.Sampler
    41  	llama2.BuildSampler(&sampler, transformer.Config.VocabSize, temperature, topp, rngSeed)
    42  
    43  	// run!
    44  	llama2.Generate(&transformer, &tokenizer, &sampler, prompt, steps)
    45  
    46  	// memory and file handles cleanup
    47  	llama2.FreeSampler(&sampler)
    48  	llama2.FreeTokenizer(&tokenizer)
    49  	llama2.FreeTransformer(&transformer)
    50  }