pedrocas15/RPC-Chat

1import gradio as gr2from rpc import load_index, generate3 4def setup():5    """ Downloads the memory mapped vector index (~10GB), installs NGT and loads the index"""6    import os7    import requests8    import subprocess9    from stream_unzip import stream_unzip10    11    DATASET_URL = os.environ.get("DATASET_URL")12    INDEX_TYPE  = os.environ.get("INDEX_TYPE")13    14    if not DATASET_URL:15        raise ValueError("DATASET_URL must be set in the environment")16        17    extract_dir = "/dev/shm/rpc-vecdb"18    os.makedirs(extract_dir, exist_ok=True)19    response = requests.get(DATASET_URL, stream=True)20    response.raise_for_status()21    22    print("Starting streaming extraction to /dev/shm...")23    for filename, filesize, file_iter in stream_unzip(response.iter_content(chunk_size=8192)):24        if isinstance(filename, bytes):25            filename = filename.decode('utf-8')26        file_path = os.path.join(extract_dir, filename)27        os.makedirs(os.path.dirname(file_path), exist_ok=True)28        with open(file_path, 'wb') as f_out:29            for chunk in file_iter:30                f_out.write(chunk)31        print(f"Extracted: {filename} -> {file_path}")32    33    files = os.listdir(extract_dir)34    files = [f for f in files if os.path.isfile(os.path.join(extract_dir, f))]35    for f in files: print(f)36    print("Index extracted")37 38    if INDEX_TYPE == "ngt":39        print("Installing NGT...")40        subprocess.check_call(["bash", "install_ngt.sh"])41        print("NGT installed")42 43    print("Loading index...")44    if INDEX_TYPE == "ngt":45        index_dir = extract_dir + "/index"46    else:47        index_dir = extract_dir48    load_index(index_path=index_dir, idx_type=INDEX_TYPE)49    print("Index loaded")50    51 52 53def respond(54    message,55    history: list[tuple[str, str]],56    user_name,57    ai_name,58    use_rpc,59    max_tokens,60    temperature,61):62    prompt = "<s>"63    for m in history:64        prompt += f"{user_name}: {m[0].strip()}\n{ai_name}: {m[1].strip()}\n"65    prompt += f"{user_name}: {message.strip()}\n{ai_name}:"66 67    response = ""68    for tok in generate(prompt, use_rpc=use_rpc, max_tokens=max_tokens):69        response += tok70        yield response71    print(history, message, response)72 73 74 75demo = gr.ChatInterface(76    respond,77    additional_inputs=[78        gr.Textbox(value="Jake", label="User name"),79        gr.Textbox(value="Sarah", label="AI name"),80        gr.Checkbox(81            label="Use RPC",82            info="Compare Normal vs. RPC-Enhanced Model",83            value=True84        ),85        gr.Slider(minimum=1, maximum=320, value=128, step=1, label="Max new tokens"),86        gr.Slider(minimum=0.1, maximum=3.0, value=0.2, step=0.1, label="Temperature (only used without RPC)"),87    ],88    description="Remember that you are talking with a 5M parameter model trained on allenai/soda, not ChatGPT"89)90 91 92if __name__ == "__main__":93    setup()94    demo.launch()95