feat: 加入GPT 模型微调功能

ChuanhuChatbot.py

@@ -5,6 +5,7 @@ logging.basicConfig(
     format="%(asctime)s [%(levelname)s] [%(filename)s:%(lineno)d] %(message)s",
 )
 
+import colorama
 import gradio as gr
 
 from modules import config
@@ -15,6 +16,7 @@ from modules.overwrites import *
 from modules.webui import *
 from modules.repo import *
 from modules.models.models import get_model
+from modules.train_func import handle_dataset_selection, handle_dataset_clear, upload_to_openai, start_training, get_training_status, add_to_models
 
 logging.getLogger("httpx").setLevel(logging.WARNING)
 
@@ -34,6 +36,7 @@ with gr.Blocks(theme=small_and_beautiful_theme) as demo:
     assert type(my_api_key)==str
     user_api_key = gr.State(my_api_key)
     current_model = gr.State(create_new_model)
+    openai_ft_file_id = gr.State("")
 
     topic = gr.State(i18n("未命名对话历史记录"))
     
@@ -188,14 +191,17 @@ with gr.Blocks(theme=small_and_beautiful_theme) as demo:
                 with gr.Tab(label=i18n("训练")):
                     with gr.Column(variant="panel"):
                         dataset_preview_json = gr.JSON(label=i18n("数据集预览"), readonly=True)
-                        upload_dataset_btn = gr.UploadButton(label = i18n("上传数据集"), file_types=[".xlsx", ".jsonl"])
+                        dataset_selection = gr.Files(label = i18n("选择数据集"), file_types=[".xlsx", ".jsonl"], file_count="single")
+                        upload_to_openai_btn = gr.Button(i18n("上传到OpenAI"), interactive=False)
+
                     with gr.Column(variant="panel"):
+                        openai_ft_suffix = gr.Textbox(label=i18n("模型名称后缀"), value="", lines=1, placeholder=i18n("可选，用于区分不同的模型"))
                         openai_train_epoch_slider = gr.Slider(label=i18n("训练轮数"), minimum=1, maximum=100, value=3, step=1, interactive=True)
                         openai_start_train_btn = gr.Button(i18n("开始训练"))
                     with gr.Column(variant="panel"):
                         openai_train_status = gr.Markdown(label=i18n("训练状态"), value=i18n("未开始训练"))
                         openai_status_refresh_btn = gr.Button(i18n("刷新状态"))
-                        add_to_models_btn = gr.Button(i18n("添加到模型列表"), interactive=False)
+                        add_to_models_btn = gr.Button(i18n("添加训练好的模型到模型列表"), interactive=False)
 
                 with gr.Tab(label=i18n("高级")):
                     gr.HTML(get_html("appearance_switcher.html").format(label=i18n("切换亮暗色主题")), elem_classes="insert-block")
@@ -485,6 +491,14 @@ with gr.Blocks(theme=small_and_beautiful_theme) as demo:
     historyFileSelectDropdown.change(**load_history_from_file_args)
     downloadFile.change(upload_chat_history, [current_model, downloadFile, user_name], [saveFileName, systemPromptTxt, chatbot])
 
+    # Train
+    dataset_selection.upload(handle_dataset_selection, dataset_selection, [dataset_preview_json, upload_to_openai_btn, status_display])
+    dataset_selection.clear(handle_dataset_clear, [], [dataset_preview_json, upload_to_openai_btn])
+    upload_to_openai_btn.click(upload_to_openai, [dataset_selection], [openai_ft_file_id, status_display], show_progress=True)
+    openai_start_train_btn.click(start_training, [openai_ft_file_id, openai_ft_suffix, openai_train_epoch_slider], [openai_train_status])
+    openai_status_refresh_btn.click(get_training_status, [], [openai_train_status, add_to_models_btn])
+    add_to_models_btn.click(add_to_models, [], [model_select_dropdown, status_display], show_progress=True)
+
     # Advanced
     max_context_length_slider.change(set_token_upper_limit, [current_model, max_context_length_slider], None)
     temperature_slider.change(set_temperature, [current_model, temperature_slider], None)

modules/index_func.py

@@ -1,7 +1,7 @@
 import os
 import logging
 
-import colorama
+import hashlib
 import PyPDF2
 from tqdm import tqdm
 
@@ -10,19 +10,6 @@ from modules.utils import *
 from modules.config import local_embedding
 
 
-def get_index_name(file_src):
-    file_paths = [x.name for x in file_src]
-    file_paths.sort(key=lambda x: os.path.basename(x))
-
-    md5_hash = hashlib.md5()
-    for file_path in file_paths:
-        with open(file_path, "rb") as f:
-            while chunk := f.read(8192):
-                md5_hash.update(chunk)
-
-    return md5_hash.hexdigest()
-
-
 def get_documents(file_src):
     from langchain.schema import Document
     from langchain.text_splitter import TokenTextSplitter
@@ -113,7 +100,7 @@ def construct_index(
     embedding_limit = None if embedding_limit == 0 else embedding_limit
     separator = " " if separator == "" else separator
 
-    index_name = get_index_name(file_src)
+    index_name = get_file_hash(file_src)
     index_path = f"./index/{index_name}"
     if local_embedding:
         from langchain.embeddings.huggingface import HuggingFaceEmbeddings

modules/train_func.py

@@ -0,0 +1,116 @@
+import os
+import logging
+import traceback
+
+import openai
+import gradio as gr
+import ujson as json
+
+import modules.presets as presets
+from modules.utils import get_file_hash
+from modules.presets import i18n
+
+def excel_to_jsonl(filepath, preview=False):
+    jsonl = []
+    with open(filepath, "rb") as f:
+        import pandas as pd
+        df = pd.read_excel(f)
+        for row in df.iterrows():
+            jsonl.append(row[1].to_dict())
+            if preview:
+                break
+    return jsonl
+
+def jsonl_save_to_disk(jsonl, filepath):
+    file_hash = get_file_hash(file_paths = [filepath])
+    os.makedirs("files", exist_ok=True)
+    save_path = f"files/{file_hash}.jsonl"
+    with open(save_path, "w") as f:
+        f.write("\n".join([json.dumps(i, ensure_ascii=False) for i in jsonl]))
+    return save_path
+
+def handle_dataset_selection(file_src):
+    logging.info(f"Loading dataset {file_src.name}...")
+    preview = ""
+    if file_src.name.endswith(".jsonl"):
+        with open(file_src.name, "r") as f:
+            preview = f.readline()
+    else:
+        preview = excel_to_jsonl(file_src.name)[0]
+    return preview, gr.update(interactive=True), "预估数据集 token 数量: 这个功能还没实现"
+
+def upload_to_openai(file_src):
+    openai.api_key = os.getenv("OPENAI_API_KEY")
+    dspath = file_src.name
+    msg = ""
+    logging.info(f"Uploading dataset {dspath}...")
+    if dspath.endswith(".xlsx"):
+        jsonl = excel_to_jsonl(dspath)
+        tmp_jsonl = []
+        for i in jsonl:
+            if "提问" in i and "答案" in i:
+                if "系统" in i :
+                    tmp_jsonl.append({
+                        "messages":[
+                            {"role": "system", "content": i["系统"]},
+                            {"role": "user", "content": i["提问"]},
+                            {"role": "assistant", "content": i["答案"]}
+                        ]
+                    })
+                else:
+                    tmp_jsonl.append({
+                        "messages":[
+                            {"role": "user", "content": i["提问"]},
+                            {"role": "assistant", "content": i["答案"]}
+                        ]
+                    })
+            else:
+                logging.warning(f"跳过一行数据，因为没有找到提问和答案: {i}")
+        jsonl = tmp_jsonl
+        dspath = jsonl_save_to_disk(jsonl, dspath)
+    try:
+        uploaded = openai.File.create(
+            file=open(dspath, "rb"),
+            purpose='fine-tune'
+            )
+        return uploaded.id, f"上传成功，文件ID: {uploaded.id}"
+    except Exception as e:
+        traceback.print_exc()
+        return "", f"上传失败，原因：{ e }"
+
+def build_event_description(id, status, trained_tokens, name=i18n("暂时未知")):
+    # convert to markdown
+    return f"""
+    #### 训练任务 {id}
+
+    模型名称：{name}
+
+    状态：{status}
+
+    已经训练了 {trained_tokens} 个token
+    """
+
+def start_training(file_id, suffix, epochs):
+    openai.api_key = os.getenv("OPENAI_API_KEY")
+    try:
+        job = openai.FineTuningJob.create(training_file=file_id, model="gpt-3.5-turbo", suffix=suffix, hyperparameters={"n_epochs": epochs})
+        return build_event_description(job.id, job.status, job.trained_tokens)
+    except Exception as e:
+        traceback.print_exc()
+        if "is not ready" in str(e):
+            return "训练出错，因为文件还没准备好。OpenAI 需要一点时间准备文件，过几分钟再来试试。"
+        return f"训练失败，原因：{ e }"
+
+def get_training_status():
+    openai.api_key = os.getenv("OPENAI_API_KEY")
+    active_jobs = [build_event_description(job["id"], job["status"], job["trained_tokens"], job["fine_tuned_model"]) for job in openai.FineTuningJob.list(limit=10)["data"] if job["status"] != "cancelled"]
+    return "\n\n".join(active_jobs), gr.update(interactive=True) if len(active_jobs) > 0 else gr.update(interactive=False)
+
+def handle_dataset_clear():
+    return gr.update(value=None), gr.update(interactive=False)
+
+def add_to_models():
+    openai.api_key = os.getenv("OPENAI_API_KEY")
+    succeeded_jobs = [job for job in openai.FineTuningJob.list(limit=10)["data"] if job["status"] == "succeeded"]
+    presets.MODELS.extend([job["fine_tuned_model"] for job in succeeded_jobs])
+    return gr.update(choices=presets.MODELS), f"成功添加了 {len(succeeded_jobs)} 个模型。"

modules/utils.py

@@ -5,14 +5,11 @@ import logging
 import commentjson as json
 import os
 import datetime
-from datetime import timezone
-import hashlib
 import csv
 import requests
 import re
 import html
-import sys
-import subprocess
+import hashlib
 
 import gradio as gr
 from pypinyin import lazy_pinyin
@@ -241,7 +238,7 @@ def convert_bot_before_marked(chat_message):
         code_block_pattern = re.compile(r"```(.*?)(?:```|$)", re.DOTALL)
         code_blocks = code_block_pattern.findall(chat_message)
         non_code_parts = code_block_pattern.split(chat_message)[::2]
-        result = []  
+        result = []
         for non_code, code in zip(non_code_parts, code_blocks + [""]):
             if non_code.strip():
                 result.append(non_code)
@@ -670,3 +667,16 @@ def auth_from_conf(username, password):
         return False
     except:
         return False
+
+def get_file_hash(file_src=None, file_paths=None):
+    if file_src:
+        file_paths = [x.name for x in file_src]
+    file_paths.sort(key=lambda x: os.path.basename(x))
+
+    md5_hash = hashlib.md5()
+    for file_path in file_paths:
+        with open(file_path, "rb") as f:
+            while chunk := f.read(8192):
+                md5_hash.update(chunk)
+
+    return md5_hash.hexdigest()

requirements.txt

@@ -21,7 +21,8 @@ duckduckgo-search
 arxiv
 wikipedia
 google.generativeai
-openai
+openai>=0.27.9
 unstructured
 google-api-python-client
 tabulate
+ujson