cyberes / sillytavern-chat-to-txt-gpt4.py

#!/usr/bin/env python3 import argparse import re from pathlib import Path import sys import json import openai import tiktoken import threading import traceback """ Convert SillyTavern jsonl chats to TXT files using AI. HOW TO USE: 1. `pip install tiktoken openai` 2. Find the chat file you want to convert. It's the `jsonl` file located in `SillyTavern/public/chats/<character name>/` 3. Run this script with `python3 sillytavern-chat-to-txt.py <path to the jsonl file> <path to where you want to save the TXT file> --key <your OpenAI API key>` This uses a temperature of 0 so don't re-run this expecting something different. If your chat is larger than the context window it will be sent in batches. After each batch, the response is written to your output file. """ class TimerThread(threading.Thread): def __init__(self, prompt:str='Waiting for response...'): super().__init__() self._stop_event = threading.Event() self.prompt = prompt def run(self): seconds = 0 while not self._stop_event.is_set(): print(f"\r{self.prompt} {seconds}s", end="") seconds += 1 self._stop_event.wait(1) def stop(self): self._stop_event.set() print('') def count_tokens(string: str, encoding_name: str = 'cl100k_base', encoding_for_model: str = None) -> int: if encoding_for_model: enc = tiktoken.encoding_for_model(encoding_for_model) else: enc = tiktoken.get_encoding(encoding_name) num_tokens = len(enc.encode(string)) return num_tokens def array_of_dicts_to_jsonl(array_of_dicts): jsonl_string = "\n".join(json.dumps(d) for d in array_of_dicts) return jsonl_string def send_to_openai(msg, model): user_msg = f"I have a jsonl transcript of an internet roleplay session. I need you to strip everything that isn't important to the story and write a summary of each message. For each message, use the format:\n\n```\n<character name>:\n<what the character says and important actions. don't use double linebreaks except for separating characters>```\n\nSome messages include things that other characters say so please organize it accordingly.\n\n{msg}" timer_thread = TimerThread(prompt=f'Sending {count_tokens(user_msg)} tokens to the AI...') timer_thread.start() try: response = openai.ChatCompletion.create( model=model, messages=[ {"role": "user", "content": user_msg} ], temperature=0, ) except Exception as e: print('Exception:', e) sys.exit(1) timer_thread.stop() return response def main(): parser = argparse.ArgumentParser(description='Convert SillyTavern jsonl files to TXT files using AI for importing into the infinite context server.') parser.add_argument('filepath', help='The path to the jsonl file to parse') parser.add_argument('output_txt', help='The output TXT file to create.') parser.add_argument('--key', required=True, help='Your OpenAI API key') parser.add_argument('--model', default='gpt-4', help='Name of the OpenAI model to use. GPT-4 seems to work the best for this. Default: gpt-4') args = parser.parse_args() openai.api_key = args.key input_jsonl = Path(args.filepath).expanduser().absolute().resolve() output_txt = Path(args.output_txt).expanduser().absolute().resolve() # Empty the file since we append to it output_txt.unlink() output_txt.touch() print('Converting chat:', input_jsonl) print('Using model:', args.model) if not input_jsonl.exists(): print('Input file does not exist:', input_jsonl) sys.exit(1) if not output_txt.parent.exists(): print('Output parent directory does not exist:', output_txt.parent) sys.exit(1) if args.model == "gpt-3.5-turbo" or args.model == "text-davinci-003": max_tokens = 3050 / 2 # div. by 2 since input+output tokens elif args.model == "gpt-4": max_tokens = 8050 / 2 else: print('Unknown model:', args.model) sys.exit(1) chatlines = [] total_tokens = 0 raw = input_jsonl.read_text().splitlines() for i in range(len(raw)): try: tmp = json.loads(raw[i]) # We don't want metadata messages if 'mes' not in tmp.keys(): continue # Trim the message down to save tokens msg = json.dumps({'name': tmp['name'], 'mes': tmp['mes']}) # We can't split messages so if one is larger than the context limit we have to quit token_count = count_tokens(msg) total_tokens += token_count if token_count > max_tokens: print('Message on line', i + 1, 'is too long at', 'token_count', 'tokens. Max tokens is', max_tokens, 'You need to decide out how to handle this.') sys.exit(1) chatlines.append(json.loads(msg)) except json.decoder.JSONDecodeError: print(f'JSON decode error on line {i + 1}:') sys.exit(1) num_chat_messages = len(chatlines) print('Total tokens:', total_tokens) while len(chatlines): ai_input_data = [] output_data = [] while True: # Check that the message fits in the max tokens ai_input_data.append(chatlines[0]) ai_input = array_of_dicts_to_jsonl(ai_input_data) token_count = count_tokens(ai_input) if token_count <= max_tokens: # Only remove a message if we fit it in the context del chatlines[0] else: # If the message is at the max token count, remove the item we just put in (we will process it next iteration) del ai_input_data[0] output_data.append(send_to_openai(ai_input, args.model).choices[0].message.content.strip()) break # If there aren't any more messages to process that means they all fit in the context if len(chatlines) == 0: output_data.append(send_to_openai(ai_input, args.model).choices[0].message.content.strip()) break if len(output_data): with open(output_txt, 'a') as f: for msg in output_data: f.write(f"{msg}\n\n\n") print(f'Converted {num_chat_messages} lines.') print('Saved to:', output_txt) if __name__ == "__main__": main()