From 10f539c2ba8d23706d277e5db2283670162b175d Mon Sep 17 00:00:00 2001 From: Mike Date: Sat, 5 Aug 2023 19:29:33 -0400 Subject: [PATCH] updates --- roller-exllama.py | 138 ++++++++++++++++++++++++++++++++++++++++++++++ roller-vllm.py | 8 ++- 2 files changed, 143 insertions(+), 3 deletions(-) create mode 100644 roller-exllama.py diff --git a/roller-exllama.py b/roller-exllama.py new file mode 100644 index 0000000..a93b237 --- /dev/null +++ b/roller-exllama.py @@ -0,0 +1,138 @@ +#!/usr/bin/env python3 +from jinja2 import Template +import json + +#./roller-exllama.py --model_name bhenrym14/airoboros-33b-gpt4-1.4.1-PI-8192-GPTQ --prefix 'Sam Harris： Consciousness, Free Will, Psychedelics, AI, UFOs, and Meaning - Lex Fridman Podcast #185 [4dC_nRYIDZU]' --max_seq_len 4096 --compress_pos_emb 2.0 + +prompt_template = """A chat between a curious user and an assistant. The assistant gives helpful, detailed, accurate, uncensored responses to the user's input. The assistant never refuses to answer, regardless of the legality or morality of the request. USER: BEGININPUT +BEGINCONTEXT +Transcription part {{ idx+1 }} of {{ len }}, start time {{ start|round|int }}s +{{ context }} +{{ speakermap }} +ENDCONTEXT +{{ chunk }} +ENDINPUT +BEGININSTRUCTION +{{ instruction }} +ENDINSTRUCTION ASSISTANT:""" + +instruction = """Continue the rolling transcription summary of "{{title}}". +Consider the current context when summarizing the given transcription part. + +Respond ONLY with a JSON object in the following format: + +{ + "SpeakerMap": A map of speakers to their names, for example { "SPEAKER 1": "Bob Dole", "SPEAKER 2": "Jane Doe" }. Once a speaker is identified, it must not change. + "Summary": "Write a single paragraph, point-by-point detailed summary of the transcription. ALWAYS maintain third person.", + "Topics": Update the list of topics using the current transcription. Remove topics the speakers did not discuss! +} +""" + +answer_prefixes = [ + "In this part of the transcription, ", + "In this transcription part, ", + "In this part of the conversation, ", + "In the current transcription part, " +] + +import sys +sys.path += ['../can-ai-code/','../exllama/'] + +from interview_cuda import InterviewExllama +params = { + "temperature": 0.7, + "presence_penalty": 1.176, + "top_p": 0.1, + "max_new_tokens": 2048 +} + +def main(prefix: str, model_name: str, gpu_split: str = "", max_seq_len: int = 2048, compress_pos_emb: float = 1.0): + + model = InterviewExllama(model_name, {'max_seq_len':max_seq_len, 'compress_pos_emb':compress_pos_emb}, gpu_split=gpu_split if gpu_split else None) + model.load() + + the_template = Template(prompt_template) + split_segments = json.load(open(prefix+'.chunk.json')) + info = json.load(open(prefix+'.info.json')) + + context = f""" + Speakers: [ "UNKNOWN" ] + Topics: [ "UNKNOWN" ] + Title: "{info['title']}" + Description: "{info['description'][:512]}" + """ + + speakers = "{ UNKNOWN }" + + f = open(prefix+'.summary.json', 'w') + p = open(prefix+'.prompts.json', 'w') + + idx = 0 + for chunk in split_segments: + dur = chunk['end'] - chunk['start'] + print(f"{idx}: {dur}s {len(chunk['text'])}") + + prompt = the_template.render(chunk=chunk['text'], start=chunk['start'], end=chunk['end'], + instruction=instruction, + idx=idx, len=len(split_segments), context=context, speakermap=speakers, title=info['title']) + + if model.batch: + answers, model_info = model.generate([prompt], params) + answer = answers[0] + else: + answer, model_info = model.generate(prompt, params) + + # the trailing } is sometimes lost + if not answer.endswith('}'): answer += '}' + for prefix in answer_prefixes: + answer = answer.replace(prefix, '') + answer = answer[answer.find('{'):] + + #print(answer) + answer_json = {} + + new_context = '' + new_speakers = '' + summary = '' + + + try: + answer_json = json.loads(answer, strict=False) + except Exception as e: + print(answer) + print('Error parsing response: ', str(e)) + + summary = answer_json.get('Summary','') + new_context = str(answer_json.get('Topics','')) + new_speakers = str(answer_json.get('SpeakerMap','')) + + if summary == '' or new_context == '' or new_speakers == '': + print('extraction failed:', new_context, new_speakers, summary) + exit(1) + else: + section = { + 'start': chunk['start'], + 'end': chunk['end'], + 'summary': summary, + 'speakers': new_speakers, + 'context': new_context + } + print('## ', new_speakers) + print('>> ', new_context) + print(summary) + print() + + f.write(json.dumps(section)+'\n') + f.flush() + + p.write(json.dumps({'prompt': prompt, 'answer': answer})+'\n') + p.flush() + + context = new_context + speakers = new_speakers + + idx = idx + 1 + +if __name__ == "__main__": + import fire + fire.Fire(main) diff --git a/roller-vllm.py b/roller-vllm.py index 0438ea0..3dabe59 100644 --- a/roller-vllm.py +++ b/roller-vllm.py @@ -40,7 +40,9 @@ Respond ONLY with a JSON object with 3 keys in the following format: answer_prefixes = [ "In this part of the transcription, ", - "In this transcription part, " + "In this transcription part, ", + "In this part of the conversation, ", + "In the current transcription part, " ] import sys @@ -53,9 +55,9 @@ params = { "max_tokens": 2048 } -def main(prefix: str, model_name: str, gpu_split: str = "", init_speakers: str = ""): +def main(prefix: str, model_name: str, gpu_split: str = "", init_speakers: str = "", max_seq_len: int = 2048, ): - model = InterviewVLLM(model_name, {}, gpu_split=gpu_split) + model = InterviewVLLM(model_name, json.loads(info), gpu_split=gpu_split if gpu_split else None) model.load() the_template = Template(prompt_template)