llama-cpp-python/examples/low_level_api/low_level_api_chat_cpp.py at main · pjq/llama-cpp-python · GitHub

History

393 lines (320 loc) · 13.2 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

335

336

337

338

339

340

341

342

343

344

345

346

347

348

349

350

351

352

353

354

355

356

357

358

359

360

361

362

363

364

365

366

367

368

369

370

371

372

373

374

375

376

377

378

379

380

381

382

383

384

385

386

387

388

389

390

391

392

393

"""

This is an example implementation of main.cpp from llama.cpp

Quirks:

* Its not exactly alike since this port is designed around programmatic I/O

* Input is always echoed if on, so it should be turned off when using "input()"

* The first antiprompt should be the userprompt like "\nUser:",

because its added when n_predict is reached (aka generation ended prematurely)

* n_predict can be set to -1 for unlimited length responses (or just a really high value)

* Instruction mode adds its own antiprompt.

You should also still be feeding the model with a "primer" prompt that

shows it the expected format.

"""

import sys

from time import time

from os import cpu_count

import llama_cpp

from common import GptParams, gpt_params_parse, gpt_random_prompt

ANSI_COLOR_RESET = "\x1b[0m"

ANSI_COLOR_YELLOW = "\x1b[33m"

ANSI_BOLD = "\x1b[1m"

ANSI_COLOR_GREEN = "\x1b[32m"

CONSOLE_COLOR_DEFAULT = ANSI_COLOR_RESET

CONSOLE_COLOR_PROMPT = ANSI_COLOR_YELLOW

CONSOLE_COLOR_USER_INPUT = ANSI_BOLD + ANSI_COLOR_GREEN

# A LLaMA interactive session

class LLaMAInteract:

def __init__(self, params: GptParams) -> None:

# input args

self.params = params

if (self.params.perplexity):

raise NotImplementedError("""************

please use the 'perplexity' tool for perplexity calculations

************""")

if (self.params.embedding):

raise NotImplementedError("""************

please use the 'embedding' tool for embedding calculations

************""")

if (self.params.n_ctx > 2048):

print(f"""warning: model does not support \

context sizes greater than 2048 tokens ({self.params.n_ctx} \

specified) expect poor results""", file=sys.stderr)

if (self.params.seed <= 0):

self.params.seed = int(time())

print(f"seed = {self.params.seed}", file=sys.stderr)

if (self.params.random_prompt):

self.params.prompt = gpt_random_prompt(self.params.seed)

# runtime args

self.input_consumed = 0

self.n_past = 0

self.first_antiprompt = []

self.remaining_tokens = self.params.n_predict

self.output_echo = self.params.input_echo

# model load

self.lparams = llama_cpp.llama_context_default_params()

self.lparams.n_ctx = self.params.n_ctx

self.lparams.n_parts = self.params.n_parts

self.lparams.seed = self.params.seed

self.lparams.memory_f16 = self.params.memory_f16

self.lparams.use_mlock = self.params.use_mlock

self.ctx = llama_cpp.llama_init_from_file(self.params.model.encode("utf8"), self.lparams)

if (not self.ctx):

raise RuntimeError(f"error: failed to load model '{self.params.model}'")

print(file=sys.stderr)

print(f"system_info: n_threads = {self.params.n_threads} / {cpu_count()} \

| {llama_cpp.llama_print_system_info().decode('utf8')}", file=sys.stderr)

# determine the required inference memory per token:

if (self.params.mem_test):

tmp = [0, 1, 2, 3]

llama_cpp.llama_eval(self.ctx, (llama_cpp.c_int * len(tmp))(*tmp), len(tmp), 0, self.n_threads)

llama_cpp.llama_print_timings(self.ctx)

self.exit()

return

# create internal context

self.n_ctx = llama_cpp.llama_n_ctx(self.ctx)

# Add a space in front of the first character to match OG llama tokenizer behavior

self.params.prompt = " " + self.params.prompt

# Load prompt file

if (self.params.file):

with open(self.params.file) as f:

self.params.prompt = f.read()

# tokenize the prompt

self.embd = []

self.embd_inp = self._tokenize(self.params.prompt)

if (len(self.embd_inp) > self.params.n_ctx - 4):

raise RuntimeError(f"error: prompt is too long ({len(self.embd_inp)} tokens, max {self.params.n_ctx - 4})")

# number of tokens to keep when resetting context

if (self.params.n_keep < 0 or self.params.n_keep > len(self.embd_inp) or self.params.instruct):

self.params.n_keep = len(self.embd_inp)

self.inp_prefix = self._tokenize(self.params.instruct_inp_prefix)

self.inp_suffix = self._tokenize(self.params.instruct_inp_suffix, False)

# in instruct mode, we inject a prefix and a suffix to each input by the user

if (self.params.instruct):

self.params.interactive_start = True

self.first_antiprompt.append(self._tokenize(self.params.instruct_inp_prefix.strip(), False))

# enable interactive mode if reverse prompt or interactive start is specified

if (len(self.params.antiprompt) != 0 or self.params.interactive_start):

self.params.interactive = True

# determine newline token

self.llama_token_newline = self._tokenize("\n", False)

if (self.params.verbose_prompt):

print(f"""

prompt: '{self.params.prompt}'

number of tokens in prompt = {len(self.embd_inp)}""", file=sys.stderr)

for i in range(len(self.embd_inp)):

print(f"{self.embd_inp[i]} -> '{llama_cpp.llama_token_to_str(self.ctx, self.embd_inp[i])}'", file=sys.stderr)

if (self.params.n_keep > 0):

print("static prompt based on n_keep: '")

for i in range(self.params.n_keep):

print(llama_cpp.llama_token_to_str(self.ctx, self.embd_inp[i]), file=sys.stderr)

print("'", file=sys.stderr)

print(file=sys.stderr)

if (self.params.interactive):

print("interactive mode on.", file=sys.stderr)

if (len(self.params.antiprompt) > 0):

for antiprompt in self.params.antiprompt:

print(f"Reverse prompt: '{antiprompt}'", file=sys.stderr)

if len(self.params.input_prefix) > 0:

print(f"Input prefix: '{self.params.input_prefix}'", file=sys.stderr)

print(f"""sampling: temp = {self.params.temp},\

top_k = {self.params.top_k},\

top_p = {self.params.top_p},\

repeat_last_n = {self.params.repeat_last_n},\

repeat_penalty = {self.params.repeat_penalty}

generate: n_ctx = {self.n_ctx}, \

n_batch = {self.params.n_batch}, \

n_predict = {self.params.n_predict}, \

n_keep = {self.params.n_keep}

""", file=sys.stderr)

# determine antiprompt tokens

for i in self.params.antiprompt:

self.first_antiprompt.append(self._tokenize(i, False))

self.last_n_tokens = [0]*self.n_ctx #TODO: deque doesnt support slices

if (params.interactive):

print("""== Running in interactive mode. ==

- Press Ctrl+C to interject at any time.

- Press Return to return control to LLaMa.

- If you want to submit another line, end your input in '\\'.

""", file=sys.stderr)

self.set_color(CONSOLE_COLOR_PROMPT)

# tokenize a prompt

def _tokenize(self, prompt, bos=True):

_arr = (llama_cpp.llama_token * (len(prompt) + 1))()

_n = llama_cpp.llama_tokenize(self.ctx, prompt.encode("utf8"), _arr, len(_arr), bos)

return _arr[:_n]

def use_antiprompt(self):

return len(self.first_antiprompt) > 0

def set_color(self, c):

if (self.params.use_color):

print(c, end="")

# generate tokens

def generate(self):

while self.remaining_tokens > 0 or self.params.interactive:

# predict

if len(self.embd) > 0:

# infinite text generation via context swapping

# if we run out of context:

# - take the n_keep first tokens from the original prompt (via n_past)

# - take half of the last (n_ctx - n_keep) tokens and recompute the logits in a batch

if (self.n_past + len(self.embd) > self.n_ctx):

n_left = self.n_past - self.params.n_keep

self.n_past = self.params.n_keep

# insert n_left/2 tokens at the start of embd from last_n_tokens

_insert = self.last_n_tokens[

self.n_ctx - int(n_left/2) - len(self.embd):-len(self.embd)

]

self.embd = _insert + self.embd

if (llama_cpp.llama_eval(

self.ctx, (llama_cpp.llama_token * len(self.embd))(*self.embd), len(self.embd), self.n_past, self.params.n_threads

) != 0):

raise Exception("Failed to llama_eval!")

self.n_past += len(self.embd)

self.embd = []

if len(self.embd_inp) <= self.input_consumed:

# out of user input, sample next token

#TODO: self.params.ignore_eos

_arr = self.last_n_tokens[-min(self.params.repeat_last_n, self.n_past):]

id = llama_cpp.llama_sample_top_p_top_k(

self.ctx,

(llama_cpp.llama_token * len(_arr))(*_arr),

len(_arr),

self.params.top_k,

self.params.top_p,

self.params.temp,

self.params.repeat_penalty,

)

self.last_n_tokens.pop(0)

self.last_n_tokens.append(id)

# replace end of text token with newline token when in interactive mode

if (id == llama_cpp.llama_token_eos() and self.params.interactive and not self.params.instruct):

id = self.llama_token_newline[0]

if (self.use_antiprompt()):

# tokenize and inject first reverse prompt

self.embd_inp += self.first_antiprompt[0]

# add it to the context

self.embd.append(id)

# echo this to console

self.output_echo = True

# decrement remaining sampling budget

self.remaining_tokens -= 1

else:

# output to console if input echo is on

self.output_echo = self.params.input_echo

# some user input remains from prompt or interaction, forward it to processing

while len(self.embd_inp) > self.input_consumed:

self.embd.append(self.embd_inp[self.input_consumed])

self.last_n_tokens.pop(0)

self.last_n_tokens.append(self.embd_inp[self.input_consumed])

self.input_consumed += 1

if len(self.embd) >= self.params.n_batch:

break

# display tokens

if self.output_echo:

for id in self.embd:

yield id

# reset color to default if we there is no pending user input

if (self.params.input_echo and len(self.embd_inp) == self.input_consumed):

self.set_color(CONSOLE_COLOR_DEFAULT)

if (self.params.interactive and len(self.embd_inp) <= self.input_consumed):

# if antiprompt is present, stop

if (self.use_antiprompt()):

if True in [

i == self.last_n_tokens[-len(i):]

for i in self.first_antiprompt

break

# if we are using instruction mode, and we have processed the initial prompt

if (self.n_past > 0 and self.params.interactive_start):

break

# end of text token

if len(self.embd) > 0 and self.embd[-1] == llama_cpp.llama_token_eos():

if (not self.params.instruct):

for i in " [end of text]\n":

yield i

break

# respect n_predict even if antiprompt is present

if (self.params.interactive and self.remaining_tokens <= 0 and self.params.n_predict != -1):

# If we arent in instruction mode, fix the current generation by appending the antiprompt.

# Makes it so if chat ends prematurely you dont append the AI's text etc.

if not self.params.instruct:

self.embd_inp += self.first_antiprompt[0]

self.n_remain = self.params.n_predict

break

self.params.interactive_start = False

def __enter__(self):

return self

def __exit__(self, type, value, tb):

self.exit()

def exit(self):

llama_cpp.llama_free(self.ctx)

self.set_color(CONSOLE_COLOR_DEFAULT)

# return past text

def past(self):

for id in self.last_n_tokens[-self.n_past:]:

yield llama_cpp.llama_token_to_str(self.ctx, id).decode("utf-8")

# write input

def input(self, prompt: str):

if (self.params.instruct and self.last_n_tokens[-len(self.inp_prefix):] != self.inp_prefix):

self.embd_inp += self.inp_prefix

self.embd_inp += self._tokenize(prompt)

if (self.params.instruct):

self.embd_inp += self.inp_suffix

# write output

def output(self):

self.remaining_tokens = self.params.n_predict

for id in self.generate():

yield llama_cpp.llama_token_to_str(self.ctx, id).decode("utf-8")

# read user input

def read_input(self):

out = ""

while (t := input()).endswith("\\"):

out += t[:-1] + "\n"

return out + t + "\n"

# interactive mode

def interact(self):

for i in self.output():

print(i,end="",flush=True)

self.params.input_echo = False

while self.params.interactive:

self.set_color(CONSOLE_COLOR_USER_INPUT)

if (self.params.instruct):

print('\n> ', end="")

self.input(self.read_input())

else:

print(self.params.input_prefix, end="")

self.input(f"{self.params.input_prefix}{self.read_input()}{self.params.output_postfix}")

print(self.params.output_postfix,end="")

self.set_color(CONSOLE_COLOR_DEFAULT)

try:

for i in self.output():

print(i,end="",flush=True)

except KeyboardInterrupt:

self.set_color(CONSOLE_COLOR_DEFAULT)

if not self.params.instruct:

print(self.params.fix_prefix,end="")

self.input(self.params.fix_prefix)

if __name__ == "__main__":

from datetime import datetime

USER_NAME="User"

AI_NAME="ChatLLaMa"

time_now = datetime.now()

prompt = f"""Text transcript of a never ending dialog, where {USER_NAME} interacts with an AI assistant named {AI_NAME}.

{AI_NAME} is helpful, kind, honest, friendly, good at writing and never fails to answer {USER_NAME}’s requests immediately and with details and precision.

There are no annotations like (30 seconds passed...) or (to himself), just what {USER_NAME} and {AI_NAME} say aloud to each other.

The dialog lasts for years, the entirety of it is shared below. It's 10000 pages long.

The transcript only includes text, it does not include markup like HTML and Markdown.

{USER_NAME}: Hello, {AI_NAME}!

{AI_NAME}: Hello {USER_NAME}! How may I help you today?

{USER_NAME}: What time is it?

{AI_NAME}: It is {time_now.strftime("%H:%M")}.

{USER_NAME}: What year is it?

{AI_NAME}: We are in {time_now.strftime("%Y")}.

{USER_NAME}: What is a cat?

{AI_NAME}: A cat is a domestic species of small carnivorous mammal. It is the only domesticated species in the family Felidae.

{USER_NAME}: Name a color.

{AI_NAME}: Blue

{USER_NAME}:"""

args = gpt_params_parse()

params = GptParams(**vars(args))

with LLaMAInteract(params) as m:

m.interact()

llama-cpp-python/examples/low_level_api/low_level_api_chat_cpp.py at main · pjq/llama-cpp-python · GitHub

Expand file tree

Breadcrumbs

low_level_api_chat_cpp.py

Latest commit

History

Breadcrumbs

low_level_api_chat_cpp.py

File metadata and controls