create_data.py

import typer
import json
from transformers import Conversation
from typing_extensions import Annotated
import httpx
import tqdm
import asyncio

app = typer.Typer()


client = httpx.AsyncClient(timeout=None)

async def run(conv: Conversation, url: str):
    payload = {"model":"tgi", "messages": conv.messages}
    response = await client.post(url, json=payload)
    content = response.json()
    message = content["choices"][0]["message"]
    message.pop("name", None)
    conv.add_message(message)




def fix_source(source):
    if source and source[0]["from"] == "gpt":
        # Skip if GPT is first to talk
        source = source[1:]
    new_source = []
    for item in source:
        role = "assistant" if item["from"] == "gpt" else "user"
        content = item["value"]
        new_source.append({"role": role, "content": content})
    return new_source


async def recreate_conversation(conversation, sem, url):
    async with sem:
        conv = Conversation()
        try:
            for message in conversation[::2]:
                assert message["role"] == "user"
                conv.add_message(message)
                await run(conv, url)
        except Exception as e:
            print(e)
            pass
        return conv.messages

@app.command()
def main(
    *,
    input_filename: Annotated[str, typer.Option("--input-filename")],
    output_filename: Annotated[str, typer.Option("--output-filename")],
    url: Annotated[str, typer.Option("--url")] = "http://localhost:8080/v1/chat/completions",
    concurrency: Annotated[int, typer.Option("--concurrency")] = 64
):
    sem = asyncio.Semaphore(concurrency)
    async def _main():
        with open(input_filename, "r") as f:
            input_data = json.loads(f.read())
        conversations = [fix_source(source["conversations"]) for source in input_data]

        futures = []
        for conversation in conversations:
            future = recreate_conversation(conversation, sem, url)
            futures.append(future)

        recreated_conversations = await tqdm.asyncio.tqdm.gather(*futures)

        with open(output_filename, "w") as f:
            json.dump(recreated_conversations, f, indent=4)
    asyncio.run(_main())


if __name__ == "__main__":
    app()