|
| 1 | +from pathlib import Path |
| 2 | +from datasets import load_dataset, concatenate_datasets |
| 3 | +import random |
| 4 | +import json |
| 5 | +import os |
| 6 | +from datetime import datetime |
| 7 | +from huggingface_hub import InferenceClient |
| 8 | +import phi_3_vision_mlx as pv |
| 9 | +import mlx.core as mx |
| 10 | +from functools import partial |
| 11 | +import fire |
| 12 | + |
| 13 | +PATH_DS = 'JosefAlbers/StampyAI-alignment-research-dataset' |
| 14 | +PROMPT_THESIS = "Based on the above bullet points, create a detailed and engaging article that explores the main themes and insights. For each bullet point, provide context, elaborate on the key ideas, and discuss their implications. Ensure the article flows logically, connects related concepts, and presents a coherent narrative." |
| 15 | +PROMPT_ANTITHESIS = "Read through the article and write a response that challenges its main ideas. Offer different viewpoints, suggest alternative explanations, and propose new approaches. Keep your response well-structured and relevant to the original content." |
| 16 | +PROMPT_SYNTHESIS = """You have an initial article and a response to it: |
| 17 | +
|
| 18 | +**Article:** |
| 19 | +{thesis} |
| 20 | +
|
| 21 | +**Response:** |
| 22 | +{antithesis} |
| 23 | +
|
| 24 | +Create an improved version of the article that incorporates insights from both the original and the response. Address conflicting ideas and present a more comprehensive view. Add new insights based on this broader perspective. Your final article should be clear, balanced, and offer a deeper understanding of the topic.""" |
| 25 | + |
| 26 | +def setup(instruction="\n<|end|>\n<|user|>\nTLDR: Summarize the following text into concise, stand-alone bullet points (max 3-5 bullet points). Each bullet point should be self-contained and provide a clear and complete idea without referencing other bullet points or the original text.", list_source=['agentmodels', 'distill', 'arbital', 'blogs', 'lesswrong', 'youtube', 'arxiv', 'special_docs'], quantize_model=False, batch_size=4, path_ds=PATH_DS): |
| 27 | + model, processor = pv.load(blind_model=True, quantize_model=quantize_model, quantize_cache=False, use_adapter=False) |
| 28 | + def aggregate(example): |
| 29 | + str_md = f"# {example['title']}\n\n{example['text']}" |
| 30 | + example['str_md'] = str_md |
| 31 | + example['len_md'] = processor(str_md)['input_ids'].size |
| 32 | + return example |
| 33 | + def summarize(example): |
| 34 | + markdowns = example['str_md'] |
| 35 | + prompts = [f'{m}{instruction}' for m in markdowns] |
| 36 | + summaries = pv.generate(prompts, preload=(model, processor), stream=False, verbose=False, max_tokens=512) |
| 37 | + example['sum_md'] = summaries |
| 38 | + return example |
| 39 | + list_ds = [] |
| 40 | + try: |
| 41 | + _ds_prev = load_dataset(path_ds, token=os.getenv("HF_WRITE_TOKEN"), split='train') |
| 42 | + list_source = [i for i in list_source if i not in _ds_prev['source']] |
| 43 | + list_ds.append(_ds_prev) |
| 44 | + except: |
| 45 | + print('Dataset not found.') |
| 46 | + for src in list_source: |
| 47 | + ds = load_dataset('StampyAI/alignment-research-dataset', src, trust_remote_code=True, split='train') |
| 48 | + ds = ds.select_columns(['id', 'source', 'title', 'text', 'url', 'date_published', 'authors', 'summary', 'source_type']) |
| 49 | + ds = ds.map(aggregate) |
| 50 | + ds = ds.filter(lambda example: 600 < example["len_md"] < 6000) |
| 51 | + if batch_size > 1: |
| 52 | + ds = ds.sort('len_md') |
| 53 | + ds = ds.map(summarize, batched=True, batch_size=batch_size) |
| 54 | + ds = ds.filter(lambda example: ('<unk>' not in example['sum_md']) and ('<|end|>' in example['sum_md'])) |
| 55 | + list_ds.append(ds) |
| 56 | + ds = concatenate_datasets(list_ds) |
| 57 | + ds.push_to_hub(path_ds, token=os.getenv("HF_WRITE_TOKEN"), private=True) |
| 58 | + |
| 59 | +def load_books(list_source=None, list_exclude=None, path_ds=PATH_DS): |
| 60 | + ds = load_dataset(path_ds, token=os.getenv("HF_READ_TOKEN", None), split='train') |
| 61 | + if list_source: |
| 62 | + list_source = [list_source] if isinstance(list_source, str) else list_source |
| 63 | + ds = ds.filter(lambda example: example['source'] in list_source) |
| 64 | + if list_exclude: |
| 65 | + list_exclude = [list_exclude] if isinstance(list_exclude, str) else list_exclude |
| 66 | + ds = ds.filter(lambda example: not any(word in example['sum_md'] for word in list_exclude)) |
| 67 | + print(f"Loaded {len(ds)} from {', '.join(set(ds['source']))}") |
| 68 | + books = ds['sum_md'] |
| 69 | + books = [i.split('\n- ') for i in books] |
| 70 | + clean_str = lambda s: s[2:] if s.startswith('- ') else s[:-7] if s.endswith('<|end|>') else s |
| 71 | + books = [[clean_str(s).strip() for s in book] for book in books] |
| 72 | + return books |
| 73 | + |
| 74 | +def pick_books(topic, list_idx, list_books, num_book=3): |
| 75 | + if topic is None: |
| 76 | + return random.sample(range(len(list_books)), num_book) |
| 77 | + list_rand = list_idx if list_idx else random.sample(range(len(list_books)), 100) |
| 78 | + list_text = [list_books[i][0] for i in list_rand] |
| 79 | + embed = pv.GteModel() |
| 80 | + l = embed(list_text) |
| 81 | + q = embed(topic) |
| 82 | + scores = mx.matmul(q, l.T) |
| 83 | + list_idx = mx.argsort(scores)[:,:-1-num_book:-1].tolist() |
| 84 | + list_idx = list_idx[0] |
| 85 | + return [list_rand[i] for i in list_idx] |
| 86 | + |
| 87 | +def get_bullets(topic='AI agents', list_source=None, list_exclude=['MIRI', 'Machine Intelligence Research Institute'], list_idx=None, num_book=3, per_book=3): |
| 88 | + books = load_books(list_source, list_exclude) |
| 89 | + list_idx = pick_books(topic, list_idx, books, num_book) |
| 90 | + print(f"Picked {list_idx}") |
| 91 | + picks = [books[i] for i in list_idx] |
| 92 | + bullets = '' |
| 93 | + for pick in picks: |
| 94 | + pick=pick[:per_book] |
| 95 | + bullets += '- ' + '\n - '.join(pick) + '\n' |
| 96 | + bullets = bullets.strip() |
| 97 | + print(f'Bullets:\n{bullets}') |
| 98 | + return bullets, list_idx |
| 99 | + |
| 100 | +def save_output(output, file_suffix=None, base_folder='syntheses'): |
| 101 | + file_suffix = f'_{file_suffix}' if file_suffix else '' |
| 102 | + os.makedirs(base_folder, exist_ok=True) |
| 103 | + date_str = datetime.now().strftime('%Y-%m-%d-%H-%M-%S') |
| 104 | + filename = os.path.join(base_folder, f'{date_str}{file_suffix}.md') |
| 105 | + with open(filename, 'w') as f: |
| 106 | + f.write(output) |
| 107 | + |
| 108 | +def synthesize(topic=None, prompt_thesis=PROMPT_THESIS, prompt_antithesis=PROMPT_ANTITHESIS, prompt_synthesis=PROMPT_SYNTHESIS, |
| 109 | + list_source=None, list_exclude=['MIRI', 'Machine Intelligence Research Institute'], |
| 110 | + list_idx=None, num_book=3, per_book=3, llm_model=None): |
| 111 | + if llm_model is None: |
| 112 | + preload = pv.load(blind_model=True, quantize_model=True) |
| 113 | + generate = partial(pv.generate, preload=preload) |
| 114 | + else: |
| 115 | + generate = partial(pv.mistral_api, api_model=llm_model, history=None, return_dict=False, verbose=False) |
| 116 | + bullets, list_idx = get_bullets(topic, list_source, list_exclude, list_idx, num_book, per_book) |
| 117 | + prompt = f"{bullets}\n\n{prompt_thesis}" |
| 118 | + thesis_output = generate(prompt) |
| 119 | + prompt_anti = f'{thesis_output}\n\n{prompt_antithesis}' |
| 120 | + antithesis_output = generate(prompt_anti) |
| 121 | + prompt_synth = prompt_synthesis.format(thesis=thesis_output, antithesis=antithesis_output) |
| 122 | + synthesis_output = generate(prompt_synth) |
| 123 | + all_output = f'Thesis:\n---\n\n{thesis_output}\n\nAntithesis:\n---\n\n{antithesis_output}\n\nSynthesis:\n---\n\n{synthesis_output}\n\nArguments:\n---\n\ndialektik.synthesize({list_source=}, {list_exclude=},{list_idx=}, {per_book=}, {llm_model=})\n\n{bullets}' |
| 124 | + save_output(all_output) |
| 125 | + return thesis_output, antithesis_output, synthesis_output |
| 126 | + |
| 127 | +if __name__ == "__main__": |
| 128 | + fire.Fire(synthesize) |
0 commit comments