app.py

# -*- encoding: utf-8 -*-
from fastapi import Path
import gradio as gr
import random
import time
from concurrent.futures import ThreadPoolExecutor
import asyncio
from functools import partial
from copy import deepcopy
from volcengine.maas import MaasService, MaasException, ChatRole
import tempfile
maas = MaasService('maas-api.ml-platform-cn-beijing.volces.com', 'cn-beijing', connection_timeout=240, socket_timeout=240)


lines=[]
cnt=0
affirm_lines=[]
def llm_infer(endpoint_id, messages, agent_network_output=True):
    req = {
        "model":{
           "endpoint_id":endpoint_id
        },
        "parameters": {
            "max_prompt_tokens": 4000,  # 最大prompt，自动截断前面的输入。0-不生效
            "max_new_tokens": 0,  # 输出文本的最大tokens限制。0-不生效
            "min_new_tokens": 0,  # 输出文本的最小tokens限制。0-不生效
            "temperature":1,  # 用于控制生成文本的随机性和创造性，Temperature值越大随机性越大。取值范围0~1
            "top_p": 0.7,  # 用于控制输出tokens的多样性，TopP值越大输出的tokens类型越丰富。取值范围0~1
            "top_k": 0,  # 选择预测值最大的k个token进行采样，取值范围0-1000。0-不生效
        },
            "messages": messages,
        }
    print(req)
    try:
        resp = maas.chat(req)
    except MaasException as e:
        print(e)
    print(resp['req_id'])
    return resp['choice']['message']['content']

def make_valid_messages(lines):
    examples = []
    for line in lines:
        messages=eval(line)

        if isinstance(messages, list):
            messages={'messages':messages}
        _messages = [messages['messages'][0]]

        pre_role = 'system'
        for part in messages['messages'][1:]:
            cur_role=part['role']
            if cur_role == 'assistant' and (pre_role == 'assistant' or pre_role=='system'):
                _messages.append({
                    'role':'user',
                    'content':""
                })
            if cur_role == 'user' and pre_role == 'user':
                _messages.append({
                    'role':'assistant',
                    'content':""
                })
            _messages.append(part)
            pre_role=cur_role
        if pre_role == 'user':
            _messages.append({
                    'role':'assistant',
                    'content':""
                })
        examples.append(
                {
                    "messages":deepcopy(_messages)
                }
            )
    return examples

def handle_click(chatbot, fix_input, ep, ak, sk, evt:gr.SelectData):
    maas.set_ak(ak)
    maas.set_sk(sk)
    global lines, cnt
    index, j = evt.index
    if fix_input is not None and len(fix_input.strip()) !=0:
        
        chatbot[index][j]=fix_input
        lines[cnt]['messages'][index + 1 + j]['content'] = fix_input
        return chatbot, ""
        
    if j == 0 or len(ep.strip())== 0 or len(ak.strip()) == 0  or len(sk.strip()) == 0:
        return chatbot
    
    messages = [
        {
        'role':'system',
        'content':lines[cnt]['messages'][0]['content']
        },
    ]
    
    print(index)
    for part in chatbot[:index]:
        messages.extend(
            [
                {
                    'role':'user',
                    'content':part[0]
                },
                {
                    'role':'assistant',
                    'content':part[1]
                }
            ]
        )
    messages.append({
        'role':'user',
        'content':chatbot[index][0]
    })
    ret = llm_infer(ep, messages)
    lines[cnt]['messages'][index + 1 + 1]['content'] = ret
    return chatbot[:index] + [[chatbot[index][0], ret]] + chatbot[index + 1:], ""

def select_line(cnt):
    global lines
    history = []
    sp=lines[cnt]['messages'][0]['content']
    for part in lines[cnt]['messages'][1:]:
        if part['role'] == 'user':
            history.append([
                part['content']
            ])
        else:
             history[-1].append(
                part['content']
            )
    return sp, history


def add_cnt():
    global cnt, lines
    cnt += 1 if cnt != len(lines) - 1 else 0
    sp, history = select_line(cnt)
    return cnt, sp, history

def sub_cnt():
    global cnt, lines
    cnt -= 1 if cnt != 0 else 0
    sp, history = select_line(cnt)
    return cnt, sp, history


def handle_file(file: gr.File):
    global lines, cnt
    lines=open(file).readlines()
    lines=make_valid_messages(lines)
    cnt=0
    sp, history=select_line(cnt)
    return history, gr.TextArea(value=sp, max_lines=5, lines=3,  interactive=True, visible=True),gr.Dropdown(choices=list(range(len(lines))), interactive=True, visible=True)
def start_eval():
    pass


def drop_download_select_line(dropdown_value):
    global cnt
    cnt = dropdown_value
    sp, history = select_line(cnt)
    return sp, history
    
def affirm_sample():
    global cnt,affirm_lines
    affirm_lines.append(lines[cnt])

    cnt += 1 if len(lines) - 1 != cnt else 0
    sp, history = select_line(cnt)

    return cnt, sp, history

def download_file():
    # https://www.gradio.app/docs/gradio/downloadbutton#demos
    global affirm_lines
    import jsonlines
    import pathlib
    ff = tempfile.NamedTemporaryFile()
    ff.name="output.jsonl"
    with jsonlines.open(ff.name, 'w') as fp:
        for line in affirm_lines:
            fp.write(line)
    
    return pathlib.Path(ff.name)

visible = gr.State(False)
sample_choices = gr.State([])
with gr.Blocks(fill_height=True) as demo:
    with gr.Row(equal_height=True): 
        with gr.Column(scale=2):
            chatbot = gr.Chatbot()
            chatbot.height=620
            affirm_btn=gr.Button(value="确认")

        with gr.Column():
            file=gr.File() # 增加些描述
            sp=gr.TextArea(visible=False)
            ak=gr.TextArea(label='ak', lines=1, max_lines=1)
            sk=gr.TextArea(label='sk', lines=1, max_lines=1)
            ep=gr.TextArea(label='ep', lines=1, max_lines=1)
            fix_input = gr.TextArea(label='手动赋值', lines=1, max_lines=3)
            sample_select=gr.Dropdown(choices=[], interactive=True, visible=False)
            with gr.Row():
                with gr.Column():
                    prev_btn = gr.Button(value='上一条')
                with gr.Column():
                    next_btn = gr.Button(value='下一条')
            eval_btn=gr.Button(value="开始生成")
            download=gr.DownloadButton("Download")


    file.upload(handle_file, [file], [chatbot, sp, sample_select])
    chatbot.select(handle_click, [chatbot, fix_input, ep, ak, sk],[chatbot, fix_input],show_progress='hidden')

    #affirm_btn=gr.Button(value="确认")
    prev_btn.click(sub_cnt,  outputs=[sample_select, sp, chatbot])
    next_btn.click(add_cnt,  outputs=[sample_select, sp, chatbot])
    affirm_btn.click(affirm_sample, inputs=[], outputs=[sample_select, sp, chatbot]) 
    download.click(download_file, inputs=[], outputs=[download])
    sample_select.change(drop_download_select_line, [sample_select], [sp, chatbot])
demo.launch()