generate-answers

#!/usr/bin/env php
<?php
/* Copyright 2024 Romain "Artefact2" Dal Maso <romain.dalmaso@artefact2.com>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *         http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

require __DIR__.'/common.php';

if($argc < 5 || $argc > 6) {
	fprintf(STDERR, "Usage: %s <model.gguf> <instruct-prefix> <instruct-suffix> <stop-sequences> [num_workers]\n", $argv[0]);
	die(1);
}
if($argc === 5) $argv[] = "1";
list(, $model_path, $prefix, $suffix, $stops, $workers) = $argv;

function spawn_server(string $model_path): int {
	if(($pid = pcntl_fork()) > 0) {
		do {
			sleep(2);
			$status = @json_decode(file_get_contents('http://127.0.0.1:24498/health'), true);
		} while(($status['status'] ?? "error") !== "ok");
		return $pid;
	}
	global $config;
	$server_args = $config['server_args']($model_path);
	passthru(escapeshellcmd(__DIR__.'/llama.cpp/server').' --host 127.0.0.1 --port 24498 -m '.escapeshellarg($model_path).' '.$server_args);
	die(0);
}

function insert_answer(int $model_id, int $prompt_id, string $answer, int $retry = 5): bool {
	/* get a new handle just for writing this answer. sqlite does not like
	 * long lived connections for writing and often gets stuck in locked
	 * states */
	$dbw = get_db();
	$ins_stmt = $dbw->prepare('INSERT INTO answers(prompt_id, model_id, answer) VALUES(:pid, :model, :answer);');
	$ins_stmt->bindValue(':model', $model_id);
	$ins_stmt->bindValue(':pid', $prompt_id);
	$ins_stmt->bindValue(':answer', $answer);
	$res = $ins_stmt->execute();
	$ins_stmt->close();
	$dbw->close();
	if($res !== false) return true;
	if($retry === 0) return false;
	sleep(1);
	return insert_answer($model, $prompt_id, $answer, $retry - 1);
}

function spawn_worker(int $model_id, int $worker_id, int $worker_count, string $prefix, string $suffix, string $stops): int {
	if(($pid = pcntl_fork()) > 0) return $pid;
	$db = $db = get_db(\SQLITE3_OPEN_READONLY);
	$pq = $db->prepare('SELECT prompts.prompt_id, prompt FROM prompts
	LEFT JOIN answers ON answers.prompt_id = prompts.prompt_id AND answers.model_id = :model
	WHERE MOD(prompts.prompt_id, :n) = :k AND answers.prompt_id IS NULL;');
	$pq->bindValue(':model', $model_id);
	$pq->bindValue(':k', $worker_id);
	$pq->bindValue(':n', $worker_count);
	$prompts = $pq->execute();
	while($row = $prompts->fetchArray()) {
		printf("\e[1m\e[33mworker %d: generating prompt_id %d\e[0m\n", $worker_id, $row['prompt_id']);
		$ctx = stream_context_create([
			'http' =>  [
				'method' => 'POST',
				'timeout' => 300,
				'header' => 'Content-Type: application/json',
				'content' => json_encode([
					'prompt' => $prefix.$row['prompt'].$suffix,
					'n_predict' => 512,
					'samplers' => 'min_p',
					'min_p' => 0.05,
					'repeat_penalty' => 1.0,
					'seed' => 42,
					'stop' => $stops,
				]),
			],
		]);
		$ans = file_get_contents('http://127.0.0.1:24498/completion', false, $ctx);
		if($ans === false) {
			var_dump($row);
			die(1);
		}
		$answer = json_decode($ans, true)['content'];
		if(!is_string($answer) || trim($answer) === '') break;
		if(insert_answer($model_id, $row['prompt_id'], $answer) === false) break;
	}
	die(0);
}

$model = pathinfo($model_path, \PATHINFO_FILENAME);
$db = get_db();
if($db->exec('BEGIN;') === false) die(1);
$mq = $db->prepare('SELECT model_id FROM models WHERE model_name = :model;');
$mq->bindValue(':model', $model);
if(($mq = $mq->execute()) === false) die(1);
if($mq = $mq->fetchArray()) {
	$model_id = $mq[0];
} else {
	$mq = $db->prepare('INSERT INTO models(model_name) VALUES(:model);');
	$mq->bindValue(':model', $model);
	if($mq->execute() === false) die(1);
	$model_id = $db->lastInsertRowID();
}
if($db->exec('COMMIT;') === false) die(1);


$todo = $db->prepare('SELECT count(*) FROM prompts
LEFT JOIN answers ON answers.prompt_id = prompts.prompt_id AND answers.model_id = :mid
WHERE answers.prompt_id IS NULL');
$todo->bindValue(':mid', $model_id);
if($todo->execute()->fetchArray()[0] === 0) {
	/* don't spawn the server for nothing */
	die(0);
}
unset($mq, $todo, $db);

$server_pid = spawn_server($model_path);
$worker_pids = [];
for($i = 0; $i < $workers; ++$i) {
	$worker_pids[] = spawn_worker($model_id, $i, $workers, $prefix, $suffix, $stops);
}

register_shutdown_function(function() { posix_kill(0, SIGKILL); });
pcntl_signal(SIGTERM, function() { die(0); });
pcntl_signal(SIGINT, function() { die(0); });

foreach($worker_pids as $p) {
	pcntl_waitpid($p, $status);
}
posix_kill($server_pid, SIGTERM);
pcntl_waitpid($server_pid, $status);