Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+models/latest_checkpoint.json filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -0,0 +1,38 @@
+# Build stage
+FROM golang:1.25-bullseye AS builder
+
+WORKDIR /app
+
+# Copy go.mod and go.sum first for caching
+COPY go.mod go.sum ./
+RUN go mod download
+
+# Copy the rest of the code
+COPY . .
+
+# Build the server binary
+RUN CGO_ENABLED=0 GOOS=linux go build -o server ./cmd/server/main.go
+
+# Final stage
+FROM debian:bullseye-slim
+
+WORKDIR /app
+
+# Install CA certificates for external downloads if needed
+RUN apt-get update && apt-get install -y ca-certificates && rm -rf /var/lib/apt/lists/*
+
+# Copy the binary from builder
+COPY --from=builder /app/server .
+
+# Copy the models directory for weights
+COPY ./models ./models
+
+# Set environment variables
+ENV PORT=7860
+ENV MODEL_PATH=models/latest_checkpoint.json
+
+# Expose the port
+EXPOSE 7860
+
+# Run the server
+CMD ["./server"]

README.md

@@ -1,12 +1,36 @@
 ---
-title: MicroGPT API
-emoji: 👁
-colorFrom: pink
+title: MicroGPT OpenAI API
+emoji: 🚀
+colorFrom: blue
 colorTo: indigo
-sdk: gradio
-sdk_version: 6.6.0
-app_file: app.py
+sdk: docker
+app_port: 7860
 pinned: false
+license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# MicroGPT OpenAI-Compatible API
+
+This Space hosts a Go-based inference server for the **MicroGPT** model, providing an OpenAI-compatible API.
+
+## API Endpoints
+
+- **`POST /v1/chat/completions`**: standard OpenAI chat format.
+- **`GET /v1/models`**: returns model metadata.
+
+## Local Test
+
+```bash
+curl -X POST http://localhost:7860/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "microgpt",
+    "messages": [
+      {"role": "user", "content": "Help me prioritize my day"}
+    ]
+  }'
+```
+
+## Repository
+
+Built with [MicroGPT Go Edition](https://github.com/Traves-Theberge/microgpt-tui-go).

cmd/server/main.go

@@ -0,0 +1,250 @@
+package main
+
+import (
+	"encoding/json"
+	"fmt"
+	"log"
+	"net/http"
+	"os"
+	"strings"
+	"time"
+
+	"microgpt-go/pkg/model"
+)
+
+type ChatMessage struct {
+	Role    string `json:"role"`
+	Content string `json:"content"`
+}
+
+type ChatCompletionRequest struct {
+	Model       string        `json:"model"`
+	Messages    []ChatMessage `json:"messages"`
+	Temperature float64       `json:"temperature"`
+	MaxTokens   int           `json:"max_tokens"`
+	TopP        float64       `json:"top_p"`
+	Stream      bool          `json:"stream"`
+}
+
+type ChatCompletionResponse struct {
+	ID      string `json:"id"`
+	Object  string `json:"object"`
+	Created int64  `json:"created"`
+	Model   string `json:"model"`
+	Choices []struct {
+		Message      ChatMessage `json:"message"`
+		Index        int         `json:"index"`
+		FinishReason string      `json:"finish_reason"`
+	} `json:"choices"`
+	Usage struct {
+		PromptTokens     int `json:"prompt_tokens"`
+		CompletionTokens int `json:"completion_tokens"`
+		TotalTokens      int `json:"total_tokens"`
+	} `json:"usage"`
+}
+
+var (
+	gpt       func(tokenID, posID int, keys, values [][][]*model.Value) []*model.Value
+	tokenizer model.TokenizerRuntime
+	config    model.TrainingCheckpointConfig
+	state     map[string][][]*model.Value
+)
+
+func initModel() {
+	ckptPath := os.Getenv("MODEL_PATH")
+	if ckptPath == "" {
+		ckptPath = "models/latest_checkpoint.json"
+	}
+	log.Printf("Loading model from %s...", ckptPath)
+	ckpt, err := model.LoadCheckpoint(ckptPath)
+	if err != nil {
+		log.Fatalf("Failed to load checkpoint: %v", err)
+	}
+
+	tokenizer, err = model.TokenizerFromCheckpoint(ckpt)
+	if err != nil {
+		log.Fatalf("Failed to load tokenizer: %v", err)
+	}
+
+	state = model.ImportState(ckpt.State)
+	config = ckpt.Config
+	gpt = model.BuildGPT(state, config.NLayer, config.NEmbd, config.NHead)
+	log.Println("Model loaded successfully.")
+}
+
+func handleChat(w http.ResponseWriter, r *http.Request) {
+	if r.Method != http.MethodPost {
+		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
+		return
+	}
+
+	var req ChatCompletionRequest
+	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+		http.Error(w, "Invalid request body", http.StatusBadRequest)
+		return
+	}
+
+	if req.Temperature <= 0 {
+		req.Temperature = 0.5
+	}
+	if req.TopP <= 0 {
+		req.TopP = 0.9
+	}
+	if req.MaxTokens <= 0 {
+		req.MaxTokens = 128
+	}
+
+	// Simple prompt construction from messages
+	var promptBuilder strings.Builder
+	for _, msg := range req.Messages {
+		role := "User"
+		if msg.Role == "assistant" {
+			role = "Assistant"
+		}
+		fmt.Fprintf(&promptBuilder, "%s: %s\n", role, msg.Content)
+	}
+	promptBuilder.WriteString("Assistant: ")
+	promptText := promptBuilder.String()
+
+	promptTokens := tokenizer.EncodeDoc(promptText)
+	if len(promptTokens) > config.BlockSize-1 {
+		promptTokens = promptTokens[len(promptTokens)-(config.BlockSize-1):]
+	}
+
+	keys := make([][][]*model.Value, config.NLayer)
+	values := make([][][]*model.Value, config.NLayer)
+	tokenID := tokenizer.BosID
+	pos := 0
+
+	// Process prompt tokens (pre-fill KV cache)
+	for _, nextID := range promptTokens {
+		if pos >= config.BlockSize {
+			break
+		}
+		_ = gpt(tokenID, pos, keys, values)
+		tokenID = nextID
+		pos++
+	}
+
+	// Generate response
+	completionTokens := 0
+	outTokens := make([]int, 0, req.MaxTokens)
+	recent := make([]int, 0, 64)
+	stopSeqs := []string{"\nUser:", "\nAssistant:"}
+
+	for pos < config.BlockSize && completionTokens < req.MaxTokens {
+		logits := gpt(tokenID, pos, keys, values)
+		recentSet := map[int]bool{}
+		for _, id := range recent {
+			recentSet[id] = true
+		}
+		weights := model.NextTokenWeights(logits, req.Temperature, 40, req.TopP, recentSet, 1.1)
+		tokenID = model.SampleWeighted(weights)
+
+		if tokenID == tokenizer.BosID {
+			break
+		}
+
+		outTokens = append(outTokens, tokenID)
+		recent = append(recent, tokenID)
+		if len(recent) > 64 {
+			recent = recent[len(recent)-64:]
+		}
+		completionTokens++
+		pos++
+
+		// Check for stop sequences in decoded text
+		fullText := tokenizer.DecodeTokens(outTokens)
+		stopFound := false
+		for _, stop := range stopSeqs {
+			if strings.Contains(fullText, stop) {
+				stopFound = true
+				break
+			}
+		}
+		if stopFound {
+			break
+		}
+	}
+
+	responseText := strings.TrimSpace(tokenizer.DecodeTokens(outTokens))
+	// Clean up any trailing stop sequence markers
+	for _, stop := range stopSeqs {
+		if idx := strings.Index(responseText, strings.TrimSpace(stop)); idx >= 0 {
+			responseText = responseText[:idx]
+		}
+	}
+
+	resp := ChatCompletionResponse{
+		ID:      fmt.Sprintf("chatcmpl-%d", time.Now().Unix()),
+		Object:  "chat.completion",
+		Created: time.Now().Unix(),
+		Model:   "microgpt",
+		Choices: []struct {
+			Message      ChatMessage `json:"message"`
+			Index        int         `json:"index"`
+			FinishReason string      `json:"finish_reason"`
+		}{
+			{
+				Message: ChatMessage{
+					Role:    "assistant",
+					Content: strings.TrimSpace(responseText),
+				},
+				Index:        0,
+				FinishReason: "stop",
+			},
+		},
+	}
+	resp.Usage.PromptTokens = len(promptTokens)
+	resp.Usage.CompletionTokens = completionTokens
+	resp.Usage.TotalTokens = resp.Usage.PromptTokens + resp.Usage.CompletionTokens
+
+	w.Header().Set("Content-Type", "application/json")
+	json.NewEncoder(w).Encode(resp)
+}
+
+func handleModels(w http.ResponseWriter, r *http.Request) {
+	resp := struct {
+		Object string `json:"object"`
+		Data   []struct {
+			ID      string `json:"id"`
+			Object  string `json:"object"`
+			Created int64  `json:"created"`
+			OwnedBy string `json:"owned_by"`
+		} `json:"data"`
+	}{
+		Object: "list",
+		Data: []struct {
+			ID      string `json:"id"`
+			Object  string `json:"object"`
+			Created int64  `json:"created"`
+			OwnedBy string `json:"owned_by"`
+		}{
+			{
+				ID:      "microgpt",
+				Object:  "model",
+				Created: time.Now().Unix(),
+				OwnedBy: "microgpt",
+			},
+		},
+	}
+	w.Header().Set("Content-Type", "application/json")
+	json.NewEncoder(w).Encode(resp)
+}
+
+func main() {
+	initModel()
+
+	http.HandleFunc("/v1/chat/completions", handleChat)
+	http.HandleFunc("/v1/models", handleModels)
+
+	port := os.Getenv("PORT")
+	if port == "" {
+		port = "7860" // Standard port for HF Spaces
+	}
+
+	log.Printf("Starting OpenAI-compatible server on port %s...", port)
+	if err := http.ListenAndServe(":"+port, nil); err != nil {
+		log.Fatalf("Failed to start server: %v", err)
+	}
+}

go.mod

@@ -0,0 +1,38 @@
+module microgpt-go
+
+go 1.25
+
+require (
+	github.com/charmbracelet/bubbles v1.0.0
+	github.com/charmbracelet/bubbletea v1.3.10
+	github.com/charmbracelet/harmonica v0.2.0
+	github.com/charmbracelet/lipgloss v1.1.0
+)
+
+require (
+	github.com/atotto/clipboard v0.1.4 // indirect
+	github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
+	github.com/charmbracelet/colorprofile v0.4.1 // indirect
+	github.com/charmbracelet/x/ansi v0.11.6 // indirect
+	github.com/charmbracelet/x/cellbuf v0.0.15 // indirect
+	github.com/charmbracelet/x/term v0.2.2 // indirect
+	github.com/clipperhouse/displaywidth v0.9.0 // indirect
+	github.com/clipperhouse/stringish v0.1.1 // indirect
+	github.com/clipperhouse/uax29/v2 v2.5.0 // indirect
+	github.com/dlclark/regexp2 v1.10.0 // indirect
+	github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f // indirect
+	github.com/google/uuid v1.3.0 // indirect
+	github.com/lucasb-eyer/go-colorful v1.3.0 // indirect
+	github.com/mattn/go-isatty v0.0.20 // indirect
+	github.com/mattn/go-localereader v0.0.1 // indirect
+	github.com/mattn/go-runewidth v0.0.19 // indirect
+	github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 // indirect
+	github.com/muesli/cancelreader v0.2.2 // indirect
+	github.com/muesli/termenv v0.16.0 // indirect
+	github.com/pkoukk/tiktoken-go v0.1.8 // indirect
+	github.com/pkoukk/tiktoken-go-loader v0.0.2 // indirect
+	github.com/rivo/uniseg v0.4.7 // indirect
+	github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect
+	golang.org/x/sys v0.38.0 // indirect
+	golang.org/x/text v0.3.8 // indirect
+)

go.sum

@@ -0,0 +1,66 @@
+github.com/atotto/clipboard v0.1.4 h1:EH0zSVneZPSuFR11BlR9YppQTVDbh5+16AmcJi4g1z4=
+github.com/atotto/clipboard v0.1.4/go.mod h1:ZY9tmq7sm5xIbd9bOK4onWV4S6X0u6GY7Vn0Yu86PYI=
+github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k=
+github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8=
+github.com/aymanbagabas/go-udiff v0.3.1 h1:LV+qyBQ2pqe0u42ZsUEtPiCaUoqgA9gYRDs3vj1nolY=
+github.com/aymanbagabas/go-udiff v0.3.1/go.mod h1:G0fsKmG+P6ylD0r6N/KgQD/nWzgfnl8ZBcNLgcbrw8E=
+github.com/charmbracelet/bubbles v1.0.0 h1:12J8/ak/uCZEMQ6KU7pcfwceyjLlWsDLAxB5fXonfvc=
+github.com/charmbracelet/bubbles v1.0.0/go.mod h1:9d/Zd5GdnauMI5ivUIVisuEm3ave1XwXtD1ckyV6r3E=
+github.com/charmbracelet/bubbletea v1.3.10 h1:otUDHWMMzQSB0Pkc87rm691KZ3SWa4KUlvF9nRvCICw=
+github.com/charmbracelet/bubbletea v1.3.10/go.mod h1:ORQfo0fk8U+po9VaNvnV95UPWA1BitP1E0N6xJPlHr4=
+github.com/charmbracelet/colorprofile v0.4.1 h1:a1lO03qTrSIRaK8c3JRxJDZOvhvIeSco3ej+ngLk1kk=
+github.com/charmbracelet/colorprofile v0.4.1/go.mod h1:U1d9Dljmdf9DLegaJ0nGZNJvoXAhayhmidOdcBwAvKk=
+github.com/charmbracelet/harmonica v0.2.0 h1:8NxJWRWg/bzKqqEaaeFNipOu77YR5t8aSwG4pgaUBiQ=
+github.com/charmbracelet/harmonica v0.2.0/go.mod h1:KSri/1RMQOZLbw7AHqgcBycp8pgJnQMYYT8QZRqZ1Ao=
+github.com/charmbracelet/lipgloss v1.1.0 h1:vYXsiLHVkK7fp74RkV7b2kq9+zDLoEU4MZoFqR/noCY=
+github.com/charmbracelet/lipgloss v1.1.0/go.mod h1:/6Q8FR2o+kj8rz4Dq0zQc3vYf7X+B0binUUBwA0aL30=
+github.com/charmbracelet/x/ansi v0.11.6 h1:GhV21SiDz/45W9AnV2R61xZMRri5NlLnl6CVF7ihZW8=
+github.com/charmbracelet/x/ansi v0.11.6/go.mod h1:2JNYLgQUsyqaiLovhU2Rv/pb8r6ydXKS3NIttu3VGZQ=
+github.com/charmbracelet/x/cellbuf v0.0.15 h1:ur3pZy0o6z/R7EylET877CBxaiE1Sp1GMxoFPAIztPI=
+github.com/charmbracelet/x/cellbuf v0.0.15/go.mod h1:J1YVbR7MUuEGIFPCaaZ96KDl5NoS0DAWkskup+mOY+Q=
+github.com/charmbracelet/x/exp/golden v0.0.0-20241011142426-46044092ad91 h1:payRxjMjKgx2PaCWLZ4p3ro9y97+TVLZNaRZgJwSVDQ=
+github.com/charmbracelet/x/exp/golden v0.0.0-20241011142426-46044092ad91/go.mod h1:wDlXFlCrmJ8J+swcL/MnGUuYnqgQdW9rhSD61oNMb6U=
+github.com/charmbracelet/x/term v0.2.2 h1:xVRT/S2ZcKdhhOuSP4t5cLi5o+JxklsoEObBSgfgZRk=
+github.com/charmbracelet/x/term v0.2.2/go.mod h1:kF8CY5RddLWrsgVwpw4kAa6TESp6EB5y3uxGLeCqzAI=
+github.com/clipperhouse/displaywidth v0.9.0 h1:Qb4KOhYwRiN3viMv1v/3cTBlz3AcAZX3+y9OLhMtAtA=
+github.com/clipperhouse/displaywidth v0.9.0/go.mod h1:aCAAqTlh4GIVkhQnJpbL0T/WfcrJXHcj8C0yjYcjOZA=
+github.com/clipperhouse/stringish v0.1.1 h1:+NSqMOr3GR6k1FdRhhnXrLfztGzuG+VuFDfatpWHKCs=
+github.com/clipperhouse/stringish v0.1.1/go.mod h1:v/WhFtE1q0ovMta2+m+UbpZ+2/HEXNWYXQgCt4hdOzA=
+github.com/clipperhouse/uax29/v2 v2.5.0 h1:x7T0T4eTHDONxFJsL94uKNKPHrclyFI0lm7+w94cO8U=
+github.com/clipperhouse/uax29/v2 v2.5.0/go.mod h1:Wn1g7MK6OoeDT0vL+Q0SQLDz/KpfsVRgg6W7ihQeh4g=
+github.com/dlclark/regexp2 v1.10.0 h1:+/GIL799phkJqYW+3YbOd8LCcbHzT0Pbo8zl70MHsq0=
+github.com/dlclark/regexp2 v1.10.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
+github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f h1:Y/CXytFA4m6baUTXGLOoWe4PQhGxaX0KpnayAqC48p4=
+github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f/go.mod h1:vw97MGsxSvLiUE2X8qFplwetxpGLQrlU1Q9AUEIzCaM=
+github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
+github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/lucasb-eyer/go-colorful v1.3.0 h1:2/yBRLdWBZKrf7gB40FoiKfAWYQ0lqNcbuQwVHXptag=
+github.com/lucasb-eyer/go-colorful v1.3.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0=
+github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
+github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
+github.com/mattn/go-localereader v0.0.1 h1:ygSAOl7ZXTx4RdPYinUpg6W99U8jWvWi9Ye2JC/oIi4=
+github.com/mattn/go-localereader v0.0.1/go.mod h1:8fBrzywKY7BI3czFoHkuzRoWE9C+EiG4R1k4Cjx5p88=
+github.com/mattn/go-runewidth v0.0.19 h1:v++JhqYnZuu5jSKrk9RbgF5v4CGUjqRfBm05byFGLdw=
+github.com/mattn/go-runewidth v0.0.19/go.mod h1:XBkDxAl56ILZc9knddidhrOlY5R/pDhgLpndooCuJAs=
+github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 h1:ZK8zHtRHOkbHy6Mmr5D264iyp3TiX5OmNcI5cIARiQI=
+github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6/go.mod h1:CJlz5H+gyd6CUWT45Oy4q24RdLyn7Md9Vj2/ldJBSIo=
+github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA=
+github.com/muesli/cancelreader v0.2.2/go.mod h1:3XuTXfFS2VjM+HTLZY9Ak0l6eUKfijIfMUZ4EgX0QYo=
+github.com/muesli/termenv v0.16.0 h1:S5AlUN9dENB57rsbnkPyfdGuWIlkmzJjbFf0Tf5FWUc=
+github.com/muesli/termenv v0.16.0/go.mod h1:ZRfOIKPFDYQoDFF4Olj7/QJbW60Ol/kL1pU3VfY/Cnk=
+github.com/pkoukk/tiktoken-go v0.1.8 h1:85ENo+3FpWgAACBaEUVp+lctuTcYUO7BtmfhlN/QTRo=
+github.com/pkoukk/tiktoken-go v0.1.8/go.mod h1:9NiV+i9mJKGj1rYOT+njbv+ZwA/zJxYdewGl6qVatpg=
+github.com/pkoukk/tiktoken-go-loader v0.0.2 h1:LUKws63GV3pVHwH1srkBplBv+7URgmOmhSkRxsIvsK4=
+github.com/pkoukk/tiktoken-go-loader v0.0.2/go.mod h1:4mIkYyZooFlnenDlormIo6cd5wrlUKNr97wp9nGgEKo=
+github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ=
+github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
+github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no=
+github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM=
+golang.org/x/exp v0.0.0-20231006140011-7918f672742d h1:jtJma62tbqLibJ5sFQz8bKtEM8rJBtfilJ2qTU199MI=
+golang.org/x/exp v0.0.0-20231006140011-7918f672742d/go.mod h1:ldy0pHrwJyGW56pPQzzkH36rKxoZW1tw7ZJpeKx+hdo=
+golang.org/x/sys v0.0.0-20210809222454-d867a43fc93e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.38.0 h1:3yZWxaJjBmCWXqhN1qh02AkOnCQ1poK6oF+a7xWL6Gc=
+golang.org/x/sys v0.38.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
+golang.org/x/text v0.3.8 h1:nAL+RVCQ9uMn3vJZbV+MRnydTJFPf8qqY42YiA6MrqY=
+golang.org/x/text v0.3.8/go.mod h1:E6s5w1FMmriuDzIBO73fBruAKo1PCIq6d2Q6DHfQ8WQ=

models/latest_checkpoint.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:964f39971833a67b2ec3a3cdd1376586aa3d3cc2b55cb11f8dc581c27a304720
+size 19575802

pkg/model/model.go

@@ -0,0 +1,375 @@
+package model
+
+import (
+	"encoding/json"
+	"fmt"
+	"math"
+	"math/rand"
+	"os"
+	"sort"
+)
+
+// Value represents a scalar for autograd
+type Value struct {
+	Data       float64
+	Grad       float64
+	Children   []*Value
+	LocalGrads []float64
+}
+
+func V(x float64) *Value {
+	return &Value{Data: x}
+}
+
+func Add(a, b *Value) *Value {
+	out := &Value{Data: a.Data + b.Data, Children: []*Value{a, b}, LocalGrads: []float64{1, 1}}
+	return out
+}
+
+func Sub(a, b *Value) *Value {
+	out := &Value{Data: a.Data - b.Data, Children: []*Value{a, b}, LocalGrads: []float64{1, -1}}
+	return out
+}
+
+func Mul(a, b *Value) *Value {
+	out := &Value{Data: a.Data * b.Data, Children: []*Value{a, b}, LocalGrads: []float64{b.Data, a.Data}}
+	return out
+}
+
+func Pow(a *Value, p float64) *Value {
+	out := &Value{Data: math.Pow(a.Data, p), Children: []*Value{a}, LocalGrads: []float64{p * math.Pow(a.Data, p-1)}}
+	return out
+}
+
+func Div(a, b *Value) *Value {
+	return Mul(a, Pow(b, -1))
+}
+
+func Neg(a *Value) *Value {
+	return Mul(a, V(-1))
+}
+
+func Log(a *Value) *Value {
+	out := &Value{Data: math.Log(a.Data), Children: []*Value{a}, LocalGrads: []float64{1 / a.Data}}
+	return out
+}
+
+func Exp(a *Value) *Value {
+	out := &Value{Data: math.Exp(a.Data), Children: []*Value{a}, LocalGrads: []float64{math.Exp(a.Data)}}
+	return out
+}
+
+func ReLU(a *Value) *Value {
+	val := 0.0
+	grad := 0.0
+	if a.Data > 0 {
+		val = a.Data
+		grad = 1
+	}
+	out := &Value{Data: val, Children: []*Value{a}, LocalGrads: []float64{grad}}
+	return out
+}
+
+func Backward(out *Value) {
+	topo := make([]*Value, 0)
+	visited := make(map[*Value]bool)
+	var buildTopo func(*Value)
+	buildTopo = func(v *Value) {
+		if !visited[v] {
+			visited[v] = true
+			for _, child := range v.Children {
+				buildTopo(child)
+			}
+			topo = append(topo, v)
+		}
+	}
+	buildTopo(out)
+
+	for _, v := range topo {
+		v.Grad = 0
+	}
+	out.Grad = 1
+	for i := len(topo) - 1; i >= 0; i-- {
+		v := topo[i]
+		for j, child := range v.Children {
+			child.Grad += v.LocalGrads[j] * v.Grad
+		}
+	}
+}
+
+func linear(x []*Value, w [][]*Value) []*Value {
+	nout := len(w)
+	nin := len(x)
+	out := make([]*Value, nout)
+	for i := 0; i < nout; i++ {
+		s := V(0)
+		for j := 0; j < nin; j++ {
+			s = Add(s, Mul(x[j], w[i][j]))
+		}
+		out[i] = s
+	}
+	return out
+}
+
+func softmax(logits []*Value) []*Value {
+	maxVal := -math.MaxFloat64
+	for _, l := range logits {
+		if l.Data > maxVal {
+			maxVal = l.Data
+		}
+	}
+	exps := make([]*Value, len(logits))
+	sumExp := V(0)
+	for i, l := range logits {
+		exps[i] = Exp(Sub(l, V(maxVal)))
+		sumExp = Add(sumExp, exps[i])
+	}
+	out := make([]*Value, len(logits))
+	invSum := Div(V(1), sumExp)
+	for i := range exps {
+		out[i] = Mul(exps[i], invSum)
+	}
+	return out
+}
+
+func rmsnorm(x []*Value) []*Value {
+	meanSq := V(0)
+	for _, v := range x {
+		meanSq = Add(meanSq, Pow(v, 2))
+	}
+	meanSq = Mul(V(1/float64(len(x))), meanSq)
+	invStd := Div(V(1), Pow(Add(meanSq, V(1e-6)), 0.5))
+	out := make([]*Value, len(x))
+	for i, v := range x {
+		out[i] = Mul(v, invStd)
+	}
+	return out
+}
+
+// TrainingCheckpoint structs
+type TrainingCheckpoint struct {
+	Version      int                      `json:"version"`
+	CreatedAt    string                   `json:"created_at"`
+	Config       TrainingCheckpointConfig `json:"config"`
+	Tokenization string                   `json:"tokenization,omitempty"`
+	BPEEncoding  string                   `json:"bpe_encoding,omitempty"`
+	BPETokenIDs  []int                    `json:"bpe_token_ids,omitempty"`
+	Vocab        []string                 `json:"vocab,omitempty"`
+	State        map[string][][]float64   `json:"state"`
+}
+
+type TrainingCheckpointConfig struct {
+	NLayer    int `json:"n_layer"`
+	NEmbd     int `json:"n_embd"`
+	NHead     int `json:"n_head"`
+	BlockSize int `json:"block_size"`
+}
+
+func ImportState(src map[string][][]float64) map[string][][]*Value {
+	out := make(map[string][][]*Value, len(src))
+	for name, mat := range src {
+		rows := make([][]*Value, len(mat))
+		for i, row := range mat {
+			r := make([]*Value, len(row))
+			for j, v := range row {
+				r[j] = V(v)
+			}
+			rows[i] = r
+		}
+		out[name] = rows
+	}
+	return out
+}
+
+func LoadCheckpoint(path string) (TrainingCheckpoint, error) {
+	b, err := os.ReadFile(path)
+	if err != nil {
+		return TrainingCheckpoint{}, err
+	}
+	var ckpt TrainingCheckpoint
+	if err := json.Unmarshal(b, &ckpt); err != nil {
+		return TrainingCheckpoint{}, err
+	}
+	if ckpt.Config.NLayer < 1 || ckpt.Config.NEmbd < 1 || ckpt.Config.NHead < 1 || ckpt.Config.BlockSize < 2 {
+		return TrainingCheckpoint{}, fmt.Errorf("invalid checkpoint config")
+	}
+	if ckpt.Config.NEmbd%ckpt.Config.NHead != 0 {
+		return TrainingCheckpoint{}, fmt.Errorf("invalid checkpoint: n_embd must be divisible by n_head")
+	}
+	return ckpt, nil
+}
+
+func BuildGPT(state map[string][][]*Value, nLayer, nEmbd, nHead int) func(tokenID, posID int, keys, values [][][]*Value) []*Value {
+	headDim := nEmbd / nHead
+	return func(tokenID, posID int, keys, values [][][]*Value) []*Value {
+		tokEmb := state["wte"][tokenID]
+		posEmb := state["wpe"][posID]
+		x := make([]*Value, len(tokEmb))
+		for i := range tokEmb {
+			x[i] = Add(tokEmb[i], posEmb[i])
+		}
+		x = rmsnorm(x)
+
+		for li := 0; li < nLayer; li++ {
+			xResidual := x
+			x = rmsnorm(x)
+			q := linear(x, state[fmt.Sprintf("layer%d.attn_wq", li)])
+			k := linear(x, state[fmt.Sprintf("layer%d.attn_wk", li)])
+			v := linear(x, state[fmt.Sprintf("layer%d.attn_wv", li)])
+			keys[li] = append(keys[li], k)
+			values[li] = append(values[li], v)
+
+			xAttn := make([]*Value, 0, nEmbd)
+			for h := 0; h < nHead; h++ {
+				hs := h * headDim
+				qH := q[hs : hs+headDim]
+
+				kH := make([][]*Value, len(keys[li]))
+				vH := make([][]*Value, len(values[li]))
+				for t := 0; t < len(keys[li]); t++ {
+					kH[t] = keys[li][t][hs : hs+headDim]
+					vH[t] = values[li][t][hs : hs+headDim]
+				}
+
+				attnLogits := make([]*Value, len(kH))
+				for t := 0; t < len(kH); t++ {
+					score := V(0)
+					for j := 0; j < headDim; j++ {
+						score = Add(score, Mul(qH[j], kH[t][j]))
+					}
+					attnLogits[t] = Div(score, V(math.Sqrt(float64(headDim))))
+				}
+				attnWeights := softmax(attnLogits)
+
+				headOut := make([]*Value, headDim)
+				for j := 0; j < headDim; j++ {
+					s := V(0)
+					for t := 0; t < len(vH); t++ {
+						s = Add(s, Mul(attnWeights[t], vH[t][j]))
+					}
+					headOut[j] = s
+				}
+				xAttn = append(xAttn, headOut...)
+			}
+
+			x = linear(xAttn, state[fmt.Sprintf("layer%d.attn_wo", li)])
+			for i := range x {
+				x[i] = Add(x[i], xResidual[i])
+			}
+
+			xResidual = x
+			x = rmsnorm(x)
+			x = linear(x, state[fmt.Sprintf("layer%d.mlp_fc1", li)])
+			for i := range x {
+				x[i] = ReLU(x[i])
+			}
+			x = linear(x, state[fmt.Sprintf("layer%d.mlp_fc2", li)])
+			for i := range x {
+				x[i] = Add(x[i], xResidual[i])
+			}
+		}
+
+		return linear(x, state["lm_head"])
+	}
+}
+
+// Sampling functions
+func SampleWeighted(weights []float64) int {
+	sum := 0.0
+	for _, w := range weights {
+		sum += w
+	}
+	r := rand.Float64() * sum
+	running := 0.0
+	for i, w := range weights {
+		running += w
+		if r <= running {
+			return i
+		}
+	}
+	return len(weights) - 1
+}
+
+func SoftmaxFloat(logits []float64) []float64 {
+	maxLogit := -math.MaxFloat64
+	for _, l := range logits {
+		if l > maxLogit {
+			maxLogit = l
+		}
+	}
+	sum := 0.0
+	out := make([]float64, len(logits))
+	for i, l := range logits {
+		out[i] = math.Exp(l - maxLogit)
+		sum += out[i]
+	}
+	for i := range out {
+		out[i] /= sum
+	}
+	return out
+}
+
+func NextTokenWeights(logits []*Value, temperature float64, topK int, topP float64, recent map[int]bool, repetitionPenalty float64) []float64 {
+	l := make([]float64, len(logits))
+	for i, v := range logits {
+		l[i] = v.Data
+		if recent[i] {
+			if l[i] >= 0 {
+				l[i] /= repetitionPenalty
+			} else {
+				l[i] *= repetitionPenalty
+			}
+		}
+		l[i] /= temperature
+	}
+	w := SoftmaxFloat(l)
+	if topK > 0 {
+		w = ApplyTopK(w, topK)
+	}
+	if topP > 0 && topP < 1.0 {
+		w = ApplyTopP(w, topP)
+	}
+	return w
+}
+
+func ApplyTopK(weights []float64, k int) []float64 {
+	if k >= len(weights) {
+		return weights
+	}
+	type kv struct {
+		i int
+		w float64
+	}
+	arr := make([]kv, len(weights))
+	for i, w := range weights {
+		arr[i] = kv{i, w}
+	}
+	sort.Slice(arr, func(i, j int) bool { return arr[i].w > arr[j].w })
+	out := make([]float64, len(weights))
+	for i := 0; i < k; i++ {
+		out[arr[i].i] = arr[i].w
+	}
+	return out
+}
+
+func ApplyTopP(weights []float64, p float64) []float64 {
+	type kv struct {
+		i int
+		w float64
+	}
+	arr := make([]kv, len(weights))
+	for i, w := range weights {
+		arr[i] = kv{i, w}
+	}
+	sort.Slice(arr, func(i, j int) bool { return arr[i].w > arr[j].w })
+	out := make([]float64, len(weights))
+	sum := 0.0
+	for i := 0; i < len(arr); i++ {
+		sum += arr[i].w
+		out[arr[i].i] = arr[i].w
+		if sum >= p {
+			break
+		}
+	}
+	return out
+}

pkg/model/tokenizer.go

@@ -0,0 +1,124 @@
+package model
+
+import (
+	"fmt"
+	"strings"
+
+	tiktoken "github.com/pkoukk/tiktoken-go"
+)
+
+type TokenizerRuntime struct {
+	Mode        string
+	CharToLocal map[rune]int
+	LocalToChar []rune
+	BpeEncoding string
+	Bpe         *tiktoken.Tiktoken
+	BpeToLocal  map[int]int
+	LocalToBPE  []int
+	UnkID       int
+	BosID       int
+}
+
+func (t TokenizerRuntime) VocabSize() int {
+	if t.Mode == "bpe_cl100k" {
+		return len(t.LocalToBPE) + 2
+	}
+	return len(t.LocalToChar) + 1
+}
+
+func (t TokenizerRuntime) EncodeDoc(doc string) []int {
+	if t.Mode == "bpe_cl100k" {
+		raw := t.Bpe.EncodeOrdinary(doc)
+		out := make([]int, 0, len(raw))
+		for _, id := range raw {
+			if local, ok := t.BpeToLocal[id]; ok {
+				out = append(out, local)
+			} else {
+				out = append(out, t.UnkID)
+			}
+		}
+		return out
+	}
+	out := make([]int, 0, len(doc))
+	for _, r := range doc {
+		if id, ok := t.CharToLocal[r]; ok {
+			out = append(out, id)
+		}
+	}
+	return out
+}
+
+func (t TokenizerRuntime) DecodeTokens(tokens []int) string {
+	if t.Mode == "bpe_cl100k" {
+		raw := make([]int, 0, len(tokens))
+		for _, local := range tokens {
+			if local >= 0 && local < len(t.LocalToBPE) {
+				raw = append(raw, t.LocalToBPE[local])
+			}
+		}
+		return t.Bpe.Decode(raw)
+	}
+	out := make([]rune, 0, len(tokens))
+	for _, id := range tokens {
+		if id >= 0 && id < len(t.LocalToChar) {
+			out = append(out, t.LocalToChar[id])
+		}
+	}
+	return string(out)
+}
+
+func TokenizerFromCheckpoint(ckpt TrainingCheckpoint) (TokenizerRuntime, error) {
+	if ckpt.Tokenization == "bpe_cl100k" || len(ckpt.BPETokenIDs) > 0 {
+		encName := strings.TrimSpace(ckpt.BPEEncoding)
+		if encName == "" {
+			encName = "cl100k_base"
+		}
+		enc, err := tiktoken.GetEncoding(encName)
+		if err != nil {
+			return TokenizerRuntime{}, err
+		}
+		localToBPE := append([]int(nil), ckpt.BPETokenIDs...)
+		bpeToLocal := make(map[int]int, len(localToBPE))
+		for i, id := range localToBPE {
+			bpeToLocal[id] = i
+		}
+		return TokenizerRuntime{
+			Mode:        "bpe_cl100k",
+			BpeEncoding: encName,
+			Bpe:         enc,
+			BpeToLocal:  bpeToLocal,
+			LocalToBPE:  localToBPE,
+			UnkID:       len(localToBPE),
+			BosID:       len(localToBPE) + 1,
+		}, nil
+	}
+	uchars, err := stringsToRunes(ckpt.Vocab)
+	if err != nil {
+		return TokenizerRuntime{}, err
+	}
+	if len(uchars) == 0 {
+		return TokenizerRuntime{}, fmt.Errorf("checkpoint has empty character vocab")
+	}
+	charToLocal := make(map[rune]int, len(uchars))
+	for i, r := range uchars {
+		charToLocal[r] = i
+	}
+	return TokenizerRuntime{
+		Mode:        "char",
+		CharToLocal: charToLocal,
+		LocalToChar: uchars,
+		BosID:       len(uchars),
+	}, nil
+}
+
+func stringsToRunes(ss []string) ([]rune, error) {
+	out := make([]rune, 0, len(ss))
+	for _, s := range ss {
+		r := []rune(s)
+		if len(r) != 1 {
+			return nil, fmt.Errorf("invalid vocab token %q: expected one rune", s)
+		}
+		out = append(out, r[0])
+	}
+	return out, nil
+}