Rewrite llm command line interface with golang (#76)

.github/workflows/cmake-darwin.yml

@@ -1,5 +1,6 @@
 name: macOS
 
+
 on:
   push:
     branches: [ "main" ]
@@ -13,8 +14,13 @@ jobs:
   build:
     runs-on: macos-14
     steps:
+    - name: Install Go
+      uses: actions/[email protected]
+      with:
+        go-version: '>=1.17.0'
+    - run: go version
     - uses: actions/checkout@v3
-    - name: Install openmp
+    - name: Install OpenMP
       run: brew install libomp
     - name: Configure CMake
       run: OpenMP_ROOT=$(brew --prefix)/opt/libomp cmake -B ${{github.workspace}}/build -DCMAKE_BUILD_TYPE=${{env.BUILD_TYPE}}

.github/workflows/cmake-windows.yml

@@ -13,6 +13,11 @@ jobs:
   build:
     runs-on: windows-latest
     steps:
+    - name: Install Go
+      uses: actions/[email protected]
+      with:
+        go-version: '>=1.17.0'
+    - run: go version
     - uses: actions/checkout@v3
     - name: Configure CMake
       run: cmake -B ${{github.workspace}}/build -DCMAKE_BUILD_TYPE=${{env.BUILD_TYPE}}

CMakeLists.txt

@@ -8,6 +8,9 @@ option(WITH_OPENMP "Build with OpenMP." ON)
 option(WITH_CUTLASS "build MatMul operators with CUTLASS." OFF)
 option(MKL_PREFIX "Prefix for MKL headers and libraries." "/opt/intel/mkl")
 
+set(CMAKE_CUDA_RUNTIME_LIBRARY Static)
+set(CMAKE_MSVC_RUNTIME_LIBRARY "MultiThreaded$<$<CONFIG:Debug>:Debug>")
+
 if(WITH_CUDA)
     add_definitions("-DLIBLLM_CUDA_ENABLED")
     find_package(CUDAToolkit REQUIRED)

go/chat/bilibili_index.go

@@ -0,0 +1,45 @@
+// The MIT License (MIT)
+//
+// Copyright (c) 2024 Xiaoyang Chen
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy of this software
+// and associated documentation files (the "Software"), to deal in the Software without
+// restriction, including without limitation the rights to use, copy, modify, merge, publish,
+// distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the
+// Software is furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in all copies or
+// substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING
+// BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
+// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+package chat
+
+import "github.com/ling0322/libllm/go/llm"
+
+type BilibiliIndex struct {
+}
+
+func (l *BilibiliIndex) Build(history []Message) (llm.Prompt, error) {
+	prompt := llm.NewPrompt()
+	if len(history) > 0 && history[0].Role == "system" {
+		prompt.AppendControlToken("<unk>")
+		prompt.AppendText(history[0].Content)
+	}
+
+	for _, message := range history {
+		if message.Role == "user" {
+			prompt.AppendControlToken("<|reserved_0|>")
+			prompt.AppendText(message.Content)
+			prompt.AppendControlToken("<|reserved_1|>")
+		} else if message.Role == "assistent" {
+			prompt.AppendText(message.Content)
+		}
+	}
+
+	return prompt, nil
+}

go/chat/chat.go

@@ -0,0 +1,63 @@
+// The MIT License (MIT)
+//
+// Copyright (c) 2024 Xiaoyang Chen
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy of this software
+// and associated documentation files (the "Software"), to deal in the Software without
+// restriction, including without limitation the rights to use, copy, modify, merge, publish,
+// distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the
+// Software is furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in all copies or
+// substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING
+// BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
+// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+package chat
+
+import (
+	"github.com/ling0322/libllm/go/llm"
+)
+
+type Message struct {
+	Role    string
+	Content string
+}
+
+type Chat struct {
+	model         llm.Model
+	promptBuilder promptBuilder
+	compConfig    llm.CompletionConfig
+}
+
+func NewChat(model llm.Model) (*Chat, error) {
+	modelName := model.GetName()
+	promptBuilder, err := newPromptBuilder(modelName)
+	if err != nil {
+		return nil, err
+	}
+
+	return &Chat{
+		model:         model,
+		promptBuilder: promptBuilder,
+		compConfig:    llm.NewCompletionConfig(),
+	}, nil
+}
+
+func (c *Chat) Chat(history []Message) (llm.Completion, error) {
+	prompt, err := c.promptBuilder.Build(history)
+	if err != nil {
+		return nil, err
+	}
+
+	comp, err := c.model.Complete(c.compConfig, prompt)
+	if err != nil {
+		return nil, err
+	}
+
+	return comp, nil
+}

go/chat/context.go

@@ -0,0 +1,31 @@
+// The MIT License (MIT)
+//
+// Copyright (c) 2024 Xiaoyang Chen
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy of this software
+// and associated documentation files (the "Software"), to deal in the Software without
+// restriction, including without limitation the rights to use, copy, modify, merge, publish,
+// distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the
+// Software is furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in all copies or
+// substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING
+// BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
+// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+package chat
+
+type QA struct {
+	Question string
+	Answer   string
+}
+
+type Context struct {
+	System   string
+	History  []QA
+	Question string
+}

go/chat/go.mod

@@ -0,0 +1,9 @@
+module github.com/ling0322/libllm/go/chat
+
+go 1.15
+
+replace github.com/ling0322/libllm/go/llm => ../llm
+
+require (
+	github.com/ling0322/libllm/go/llm v1.0.0
+)

go/chat/llama.go

@@ -0,0 +1,43 @@
+// The MIT License (MIT)
+//
+// Copyright (c) 2024 Xiaoyang Chen
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy of this software
+// and associated documentation files (the "Software"), to deal in the Software without
+// restriction, including without limitation the rights to use, copy, modify, merge, publish,
+// distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the
+// Software is furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in all copies or
+// substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING
+// BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
+// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+package chat
+
+import "github.com/ling0322/libllm/go/llm"
+
+type Llama struct {
+}
+
+func (l *Llama) Build(history []Message) (llm.Prompt, error) {
+	prompt := llm.NewPrompt()
+	prompt.AppendControlToken("<|begin_of_text|>")
+	for _, message := range history {
+		prompt.AppendControlToken("<|start_header_id|>")
+		prompt.AppendText(message.Role)
+		prompt.AppendControlToken("<|end_header_id|>")
+		prompt.AppendText("\n\n" + message.Content)
+		prompt.AppendControlToken("<|eot_id|>")
+	}
+
+	prompt.AppendControlToken("<|start_header_id|>")
+	prompt.AppendText("assistant")
+	prompt.AppendControlToken("<|end_header_id|>")
+	prompt.AppendText("\n\n")
+	return prompt, nil
+}

go/chat/prompt_builder.go

@@ -0,0 +1,40 @@
+// The MIT License (MIT)
+//
+// Copyright (c) 2024 Xiaoyang Chen
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy of this software
+// and associated documentation files (the "Software"), to deal in the Software without
+// restriction, including without limitation the rights to use, copy, modify, merge, publish,
+// distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the
+// Software is furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in all copies or
+// substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING
+// BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
+// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+package chat
+
+import (
+	"fmt"
+
+	"github.com/ling0322/libllm/go/llm"
+)
+
+type promptBuilder interface {
+	Build(history []Message) (llm.Prompt, error)
+}
+
+func newPromptBuilder(modelName string) (promptBuilder, error) {
+	if modelName == "llama" {
+		return &Llama{}, nil
+	} else if modelName == "index" {
+		return &BilibiliIndex{}, nil
+	} else {
+		return nil, fmt.Errorf("unexpected model name %s", modelName)
+	}
+}

go/cmd/go.mod

@@ -1,9 +1,12 @@
 module main
 
-go 1.22.4
+go 1.15
 
 replace github.com/ling0322/libllm/go/llm => ../llm
 
+replace github.com/ling0322/libllm/go/chat => ../chat
+
 require (
 	github.com/ling0322/libllm/go/llm v1.0.0
+	github.com/ling0322/libllm/go/chat v1.0.0
 )