app.py

from threading import Thread
import gradio as gr
import inspect
from gradio import routes
from typing import List, Type

import requests, os, re, asyncio, queue
import math
import time
import datetime
import requests, json

from pprint import pprint
import hivemind
from petals.constants import PUBLIC_INITIAL_PEERS
from health import fetch_health_state

dht = hivemind.DHT(initial_peers=PUBLIC_INITIAL_PEERS, client_mode=True, start=True)
model_name = "quantumaikr/llama-2-70b-fb16-korean"

loop = asyncio.get_event_loop()
# Monkey patch
def get_types(cls_set: List[Type], component: str):
    docset = []
    types = []
    if component == "input":
        for cls in cls_set:
            doc = inspect.getdoc(cls)
            doc_lines = doc.split("\n")
            docset.append(doc_lines[1].split(":")[-1])
            types.append(doc_lines[1].split(")")[0].split("(")[-1])
    else:
        for cls in cls_set:
            doc = inspect.getdoc(cls)
            doc_lines = doc.split("\n")
            docset.append(doc_lines[-1].split(":")[-1])
            types.append(doc_lines[-1].split(")")[0].split("(")[-1])
    return docset, types
routes.get_types = get_types

# App code

account_list = dict()

account_list['id'] = "pass"

name_list = dict()
name_list['id'] = 'name'

p2p_list = dict()
p2p_list['id'] = '11111111'

def chat(id, prompt):
    
    return "AI 응답입니다."


def register(id, pw):
    if id in account_list:
        return "exist"
    else:
        account_list[id] = pw
        return "ok"

def login(id, pw):
    if id in account_list:
        if account_list[id] == pw:
            return "ok"
        else:
            return "password error"
    else:
        return "no id"

def add_name(id, name):
    name_list[id] = name
    return "ok"

def get_name(id):
    if id in name_list:
        return name_list[id]
    else:
        return "no id"

def get_id(name):
    reverse_dict= dict(map(reversed,name_list.items()))
    if name in reverse_dict:
        return reverse_dict[name]
    else:
        return "no name"

def add_p(id, p_id):
    p2p_list[id] = p_id
    return "ok"

def get_p(id):
    if id in p2p_list:
        return p2p_list[id]
    else:
        return "no id"

def get_id_from_p2p(i):
    reverse_dict= dict(map(reversed,p2p_list.items()))
    if i in reverse_dict:
        return reverse_dict[i]
    else:
        return "no id"

# Blockchain code

def get_peers(name):
    data = fetch_health_state(dht)
    out = []
    for d in data['model_reports']:
        if d['name'] == name:
            for r in d['server_rows']:
                out.append(r['peer_id'])

    return out

with gr.Blocks() as demo:
    count = 0
    aa = gr.Interface(
      fn=chat,
      inputs=["text","text"],
      outputs="text",
      description="chat, ai 응답을 반환합니다. \n /run/predict",
    )

    rr = gr.Interface(
      fn=register,
      inputs=["text", "text"],
      outputs="text",
      description="register, 회원가입(성공시:ok, 중복시:exist 반환)\n /run/predict_1",
    )

    ll = gr.Interface(
      fn=login,
      inputs=["text", "text"],
      outputs="text",
      description="login, 로그인(성공시: ok, 실패시: password error, 아이디가 없으면: no id) \n /run/predict_2",
    )

    ad = gr.Interface(
      fn=add_name,
      inputs=["text", "text"],
      outputs="text",
      description="add_name, id로 닉네임 추가. ok 반환.\n /run/predict_3",
    )
    
    nn = gr.Interface(
      fn=get_name,
      inputs=["text"],
      outputs="text",
      description="get_name, id로 닉네임 반환(없으면 no id)\n /run/predict_4",
    )

    nnn = gr.Interface(
      fn=get_id,
      inputs=["text"],
      outputs="text",
      description="get_name, 닉네임으로 id 반환(없으면 no name)\n /run/predict_5",
    )
    
    adp = gr.Interface(
      fn=add_p,
      inputs=["text", "text"],
      outputs="text",
      description="add_p, id로 p2p id 추가. ok 반환. \n /run/predict_6",
    )
    
    nnp = gr.Interface(
      fn=get_p,
      inputs=["text"],
      outputs="text",
      description="get_p, id로 p2p id 반환. 없으면 no id. \n /run/predict_7",
    )

    nnp = gr.Interface(
      fn=get_id_from_p2p,
      inputs=["text"],
      outputs="text",
      description="get_p, p2p id로 일반 id 반환. 없으면 no id. \n /run/predict_8",
    )

    gpeer = gr.Interface(
      fn=get_peers,
      inputs=["text"],
      outputs="text",
      description="get_peers, 해당 모델을 분산처리중인 peer들의 p2p id list 반환\n /run/predict_8",
    )
    
    demo.queue(max_size=32).launch(enable_queue=True)