import argparse
import asyncio
import json
import pprint
import traceback
from typing import cast

from datasets import Dataset, DatasetDict, load_dataset
from pydantic import BaseModel
from tqdm import tqdm
from twopir import llms
from twopir.core import (
    ChatMessagesTemplate,
    ChatMessageTemplate,
    ChatRole,
    PromptTemplate,
)
from pathlib import Path
from collections import defaultdict
import time
import random
import re


def extract_base_dataset():
    ds = load_dataset("ashraq/financial-news-articles", split="train")
    ds = ds.filter(lambda ex: len(ex["text"]) > 2000)

    print(f"{ds=}")

    article = []
    for ex in ds:
        article.append("{}\n\n{}".format(ex["title"], ex["text"]))
    ds = ds.add_column("article", article)

    ds = ds.select(range(20000))
    ds.push_to_hub("withpi/financial-news-articles", private=True)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--max",
        help="max count",
        default=10,
        type=int,
    )
    args = parser.parse_args()

    extract_base_dataset()


if __name__ == "__main__":
    main()