import (
    "fmt"
    "os"

    "github.com/gomlx/compute/support/humanize"
    "github.com/gomlx/go-huggingface/hub"
    "github.com/gomlx/go-huggingface/tokenizers"
)

var (
	// HuggingFace authentication token read from environment.
	// It can be created in https://huggingface.co
	// Some files may require it for downloading.
	hfAuthToken = os.Getenv("HF_TOKEN")

	// Model IDs we use for testing.
	hfModelIDs = []string{
		"google/gemma-2-2b-it",
		"sentence-transformers/all-MiniLM-L6-v2",
		"protectai/deberta-v3-base-zeroshot-v1-onnx",
		"KnightsAnalytics/distilbert-base-uncased-finetuned-sst-2-english",
		"KnightsAnalytics/distilbert-NER",
		"SamLowe/roberta-base-go_emotions-onnx",
	}
)

for _, modelID := range hfModelIDs {
	fmt.Printf("\n%s:\n", modelID)
	repo := hub.New(modelID).WithAuth(hfAuthToken)
	for fileInfo, err := range repo.IterFileInfos() {
		if err != nil { panic(err) }
		fmt.Printf("\t%s - %s\n", fileInfo.Name, humanize.Bytes(fileInfo.Size))
	}
}

repo := hub.New("google/gemma-2-2b-it").WithAuth(hfAuthToken)
tokenizer, err := tokenizers.New(repo)
if err != nil { panic(err) }

sentence := "The book is on the table."
tokens := tokenizer.Encode(sentence)
fmt.Printf("Sentence:\t%s\n", sentence)
fmt.Printf("Tokens:  \t%v\n", tokens)

for _, modelID := range hfModelIDs {
	fmt.Printf("\n%s:\n", modelID)
	repo := hub.New(modelID).WithAuth(hfAuthToken)
	config, err := tokenizers.GetConfig(repo)
	if err != nil { panic(err) }
	fmt.Printf("\ttokenizer_class=%s\n", config.TokenizerClass)
}

tokenizer := ... // see previous example

// Start bucket runner in a separate goroutine.
var wg sync.WaitGroup
bucketsInputChan := make(chan bucket.SentenceRef)
bucketsOutputChan := make(chan bucket.Bucket, 10)
bkt := bucket.New(tokenizer).
	ByTwoBitBucketBudget(8*1024, 16).  // ~8K total tokens per bucket, ~20% padding overhead
	WithMaxParallelization(-1)
wg.Go(func() {
	bkt.Run(bucketsInputChan, bucketsOutputChan)
})
...

import dtok "github.com/daulet/tokenizers"

%%
modelID := "KnightsAnalytics/all-MiniLM-L6-v2"
repo := hub.New(modelID).WithAuth(hfAuthToken)
localFile := must.M1(repo.DownloadFile("tokenizer.json"))
tokenizer := must.M1(dtok.FromFile(localFile))
defer tokenizer.Close()
tokens, _ := tokenizer.Encode(sentence, true)

fmt.Printf("Sentence:\t%s\n", sentence)
fmt.Printf("Tokens:  \t%v\n", tokens)

import (
	"github.com/gomlx/go-huggingface/hub"
	"github.com/gomlx/go-huggingface/models/transformer"
	"github.com/gomlx/gomlx/ml/model"
)

// 1. Download configuration and weights from HuggingFace
repo := hub.New("tencent/KaLM-Embedding-Gemma3-12B-2511").WithAuth(hfAuthToken)
hfModel, err := transformer.LoadModel(repo)
if err != nil { panic(err) }
hfModel.WithCausalMask(true)
tokenizer := must1(hfModel.GetTokenizer())
padID, err := tokenizer.SpecialTokenID(api.TokPad)

// Print a summary of the model features and sizes:
fmt.Println(hfModel.Description())

// 2. Load the loaded weights to a GoMLX context
backend := compute.MustNew()
store := model.NewStore()
err := hfModel.LoadStore(backend, store)  // Load model weights into the store.

// 3. Build a GoMLX graph and executor for the model.
kalmExec, err := model.NewExec1(testBackend, testStore, func(scope *model.Scope, tokens *graph.Node) *graph.Node {
	x := hfModel.SentenceEmbeddingGraph(scope, tokens, nil)
	return graph.ConvertDType(x, dtypes.Float32)
})

// 4. Embed sentences ...
tokens := tokenzier.Encode(prompt)
embeddings := kalmExec.MustCall(tokens)

//...

import "github.com/gomlx/go-huggingface/datasets"

// Print dataset info: configurations, splits, sizes and features.
ds := datasets.New("HuggingFaceFW/fineweb").WithAuth(hfAuthToken)
fmt.Println(ds.String())

//... (some info messages)...
type FinewebRecord struct {
	Text string `json:"text" parquet:"text"`
	ID string `json:"id" parquet:"id"`
	Dump string `json:"dump" parquet:"dump"`
	URL string `json:"url" parquet:"url"`
	Date string `json:"date" parquet:"date"`
	FilePath string `json:"file_path" parquet:"file_path"`
	Language string `json:"language" parquet:"language"`
	LanguageScore float64 `json:"language_score" parquet:"language_score"`
	TokenCount int64 `json:"token_count" parquet:"token_count"`
}

var (
    FineWebID = "HuggingFaceFW/fineweb"
    FineWebConfig = "sample-10BT"
    FineWebSplit = "train"
)

// FinewwebRecord generated with github.com/gomlx/go-huggingface/cmd/generate_dataset_structs
type FinewebRecord struct {
	Text string `json:"text" parquet:"text"`
	ID string `json:"id" parquet:"id"`
	Dump string `json:"dump" parquet:"dump"`
	URL string `json:"url" parquet:"url"`
	Date string `json:"date" parquet:"date"`
	FilePath string `json:"file_path" parquet:"file_path"`
	Language string `json:"language" parquet:"language"`
	LanguageScore float64 `json:"language_score" parquet:"language_score"`
	TokenCount int64 `json:"token_count" parquet:"token_count"`
}

// TrimString returns s trimmed to at most maxLength runes. If trimmed it appends "…" at the end.
func TrimString(s string, maxLength int) string {
    if utf8.RuneCountInString(s) <= maxLength {
        return s
    }
    runes := []rune(s)
    return string(runes[:maxLength-1]) + "…"
}

%%
ds := datasets.New(FineWebID)
ds.Verbosity = 2
count := 0
const limit = 10
for row, err := range datasets.IterParquetFromDataset[FinewebRecord](ds, FineWebConfig, FineWebSplit) {
    if err != nil { panic(err) }
	fmt.Printf("Record #%02d:\tScore=%.3f Text=%q, URL=[%s]\n", count+1, row.LanguageScore, TrimString(row.Text, 50), TrimString(row.URL, 40))
    count++
    if count >= limit { break }
}
fmt.Printf("%d records read!", count)

import (
    "github.com/gomlx/compute"
    "github.com/gomlx/go-huggingface/tokenizers/bucket"
    onnxparser "github.com/gomlx/onnx-gomlx/onnx/parser"
    "github.com/gomlx/gomlx/core/graph"
    "github.com/gomlx/gomlx/ml/model"

    // Default backends.
    _ "github.com/gomlx/gomlx/backends/default"
)

%%
// Get ONNX model.
repo := hub.New("sentence-transformers/all-MiniLM-L6-v2").WithAuth(hfAuthToken)
onnxFilePath, err := repo.DownloadFile("onnx/model.onnx")
if err != nil { panic(err) }
onnxModel, err := onnxparser.FromFile(onnxFilePath)
if err != nil { panic(err) }

// Convert ONNX variables to a GoMLX store:
store := model.NewStore()
err = onnxModel.VariablesToScope(store.RootScope())
if err != nil { panic(err) }

// Tokenize sentences.
tokenizer := must.M1(tokenizers.New(repo))
sentences := []string{
    "This is an example sentence", 
    "Each sentence is converted"}
batchSize := len(sentences)
sentencesChan := make(chan bucket.SentenceRef, batchSize)
bucketChan := make(chan bucket.Bucket, 1)
bucketizer := bucket.New(tokenizer).ByTwoBitBucket(batchSize, 8)
var wg sync.WaitGroup
wg.Go(func() { bucketizer.Run(sentencesChan, bucketChan) })
wg.Go(func() { 
    for i, s := range sentences { sentencesChan <- bucket.SentenceRef{s, i} }
    close(sentencesChan)
})

// Create GoMLX model, and its executor:
miniLMExec := model.MustNewExec1(
    compute.MustNew(), store, 
    func (scope *model.Scope, tokenIDs *graph.Node) *graph.Node {
        tokenIDs = graph.Reshape(tokenIDs, batchSize, -1)
        mask := graph.LogicalNot(graph.IsZero(tokenIDs))
        return onnxModel.CallGraph(scope, tokenIDs.Graph(), map[string]*graph.Node{
            "input_ids": tokenIDs,
            "attention_mask": graph.ConvertDType(mask, dtypes.Int64),
            "token_type_ids": graph.ZerosLike(tokenIDs)})[0]
    })

// Loop over batches:
for bucket := range bucketChan {
    tokenIDs := bucket.Batch
    embeddings := miniLMExec.MustCall(tokenIDs)
    fmt.Printf("Tokens: \t%v\n", tokenIDs)
    fmt.Printf("Embeddings:\t%s\n", embeddings)
}