context_test.go 620 Bytes
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
//go:build integration

package integration

import (
	"context"
	"testing"
	"time"

	"github.com/ollama/ollama/api"
)

func TestContextExhaustion(t *testing.T) {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
14
15
	// Longer needed for small footprint GPUs
	ctx, cancel := context.WithTimeout(context.Background(), 6*time.Minute)
16
17
18
19
20
21
22
23
24
25
26
27
	defer cancel()
	// Set up the test data
	req := api.GenerateRequest{
		Model:  "llama2",
		Prompt: "Write me a story with a ton of emojis?",
		Stream: &stream,
		Options: map[string]interface{}{
			"temperature": 0,
			"seed":        123,
			"num_ctx":     128,
		},
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
28
	GenerateTestHelper(ctx, t, req, []string{"once", "upon", "lived"})
29
}