Vodalus / app.py
BeTaLabs's picture
Update app.py
75d10b9 verified
raw
history blame contribute delete
No virus
39.8 kB
import gradio as gr
from gradio import update
import json
import re
from datetime import datetime
from typing import Literal
import os
import importlib
from llm_handler import send_to_llm_wrapper
from main import generate_data, PROMPT_1
from topics import TOPICS
from system_messages import SYSTEM_MESSAGES_VODALUS
import random
from params import load_params, save_params
import pandas as pd
import csv
from datasets import load_dataset
from huggingface_hub import list_datasets, HfApi, hf_hub_download
from gradio.components import State
ANNOTATION_CONFIG_FILE = "annotation_config.json"
OUTPUT_FILE_PATH = "dataset.jsonl"
llm_provider_state = State("")
def load_llm_config():
params = load_params()
return (
params.get('PROVIDER', ''),
params.get('BASE_URL', ''),
params.get('MODEL', ''), # Add this line
params.get('WORKSPACE', ''),
params.get('API_KEY', ''),
params.get('max_tokens', 2048),
params.get('temperature', 0.7),
params.get('top_p', 0.9),
params.get('frequency_penalty', 0.0),
params.get('presence_penalty', 0.0)
)
def save_llm_config(provider, base_url, model, workspace, api_key, max_tokens, temperature, top_p, frequency_penalty, presence_penalty):
save_params({
'PROVIDER': provider,
'BASE_URL': base_url,
'MODEL': model, # Add this line
'WORKSPACE': workspace,
'API_KEY': api_key,
'max_tokens': max_tokens,
'temperature': temperature,
'top_p': top_p,
'frequency_penalty': frequency_penalty,
'presence_penalty': presence_penalty
})
return "LLM configuration saved successfully"
def update_model_visibility(provider):
return gr.update(visible=provider in ["local-model", "openai"])
def load_annotation_config():
try:
with open(ANNOTATION_CONFIG_FILE, 'r') as f:
return json.load(f)
except FileNotFoundError:
return {
"quality_scale": {
"name": "Relevance for Training",
"description": "Rate the relevance of this entry for training",
"scale": [
{"value": "1", "label": "Invalid"},
{"value": "2", "label": "Somewhat invalid"},
{"value": "3", "label": "Neutral"},
{"value": "4", "label": "Somewhat valid"},
{"value": "5", "label": "Valid"}
]
},
"tag_categories": [
{
"name": "High Quality Indicators",
"type": "multiple",
"tags": ["Well-formatted", "Informative", "Coherent", "Engaging"]
},
{
"name": "Low Quality Indicators",
"type": "multiple",
"tags": ["Poorly formatted", "Lacks context", "Repetitive", "Irrelevant"]
},
{
"name": "Content Warnings",
"type": "multiple",
"tags": ["Offensive language", "Hate speech", "Violence", "Adult content"]
}
],
"free_text_fields": [
{
"name": "Additional Notes",
"description": "Any other observations or comments"
}
]
}
def load_csv_dataset(file_path):
data = []
with open(file_path, 'r') as f:
reader = csv.DictReader(f)
for row in reader:
data.append(row)
return data
def load_txt_dataset(file_path):
with open(file_path, 'r') as f:
return [{"content": line.strip()} for line in f if line.strip()]
def save_annotation_config(config):
with open(ANNOTATION_CONFIG_FILE, 'w') as f:
json.dump(config, f, indent=2)
def load_jsonl_dataset(file_path):
if not os.path.exists(file_path):
return []
with open(file_path, 'r') as f:
return [json.loads(line.strip()) for line in f if line.strip()]
def load_dataset(file):
if file is None:
return "", 0, 0, "No file uploaded", "3", [], [], [], ""
file_path = file.name
file_extension = os.path.splitext(file_path)[1].lower()
if file_extension == '.csv':
data = load_csv_dataset(file_path)
elif file_extension == '.txt':
data = load_txt_dataset(file_path)
elif file_extension == '.jsonl':
data = load_jsonl_dataset(file_path)
else:
return "", 0, 0, f"Unsupported file type: {file_extension}", "3", [], [], [], ""
if not data:
return "", 0, 0, "No data found in the file", "3", [], [], [], ""
first_row = json.dumps(data[0], indent=2)
return first_row, 0, len(data), f"Row: 1/{len(data)}", "3", [], [], [], ""
def save_row(file_path, index, row_data):
file_extension = file_path.split('.')[-1].lower()
if file_extension == 'jsonl':
save_jsonl_row(file_path, index, row_data)
elif file_extension == 'csv':
save_csv_row(file_path, index, row_data)
elif file_extension == 'txt':
save_txt_row(file_path, index, row_data)
else:
raise ValueError(f"Unsupported file format: {file_extension}")
return f"Row {index} saved successfully"
def save_jsonl_row(file_path, index, row_data):
with open(file_path, 'r') as f:
lines = f.readlines()
lines[index] = row_data + '\n'
with open(file_path, 'w') as f:
f.writelines(lines)
def save_csv_row(file_path, index, row_data):
df = pd.read_csv(file_path)
row_dict = json.loads(row_data)
for col, value in row_dict.items():
df.at[index, col] = value
df.to_csv(file_path, index=False)
def save_txt_row(file_path, index, row_data):
with open(file_path, 'r') as f:
lines = f.readlines()
row_dict = json.loads(row_data)
lines[index] = row_dict.get('content', '') + '\n'
with open(file_path, 'w') as f:
f.writelines(lines)
def get_row(file_path, index):
data = load_jsonl_dataset(file_path)
if not data:
return "", 0
if 0 <= index < len(data):
return json.dumps(data[index], indent=2), len(data)
return "", len(data)
def json_to_markdown(json_str):
try:
data = json.loads(json_str)
markdown = f"# System\n\n{data['system']}\n\n# Instruction\n\n{data['instruction']}\n\n# Response\n\n{data['response']}"
return markdown
except json.JSONDecodeError:
return "Error: Invalid JSON format"
def markdown_to_json(markdown_str):
sections = re.split(r'#\s+(System|Instruction|Response)\s*\n', markdown_str)
if len(sections) != 7: # Should be: ['', 'System', content, 'Instruction', content, 'Response', content]
return "Error: Invalid markdown format"
json_data = {
"system": sections[2].strip(),
"instruction": sections[4].strip(),
"response": sections[6].strip()
}
return json.dumps(json_data, indent=2)
def navigate_rows(file_path: str, current_index: int, direction: Literal["prev", "next"], metadata_config):
new_index = max(0, current_index + (-1 if direction == "prev" else 1))
return load_and_show_row(file_path, new_index, metadata_config)
def load_and_show_row(file_path, index, metadata_config):
row_data, total = get_row(file_path, index)
if not row_data:
return ("", index, total, f"Row: {index + 1}/{total}", "3", [], [], [], "")
try:
data = json.loads(row_data)
except json.JSONDecodeError:
return (row_data, index, total, f"Row: {index + 1}/{total}", "3", [], [], [], "Error: Invalid JSON")
metadata = data.get("metadata", {}).get("annotation", {})
quality = metadata.get("quality", "3")
high_quality_tags = metadata.get("tags", {}).get("high_quality", [])
low_quality_tags = metadata.get("tags", {}).get("low_quality", [])
toxic_tags = metadata.get("tags", {}).get("toxic", [])
other = metadata.get("free_text", {}).get("Additional Notes", "")
return (row_data, index, total, f"Row: {index + 1}/{total}", quality,
high_quality_tags, low_quality_tags, toxic_tags, other)
def save_row_with_metadata(file_path, index, row_data, config, quality, high_quality_tags, low_quality_tags, toxic_tags, other):
data = json.loads(row_data)
metadata = {
"annotation": {
"quality": quality,
"tags": {
"high_quality": high_quality_tags,
"low_quality": low_quality_tags,
"toxic": toxic_tags
},
"free_text": {
"Additional Notes": other
}
}
}
data["metadata"] = metadata
return save_row(file_path, index, json.dumps(data))
def update_annotation_ui(config):
quality_choices = [(item["value"], item["label"]) for item in config["quality_scale"]["scale"]]
quality_label = gr.Radio(
label=config["quality_scale"]["name"],
choices=quality_choices,
info=config["quality_scale"]["description"]
)
tag_components = []
for category in config["tag_categories"]:
tag_component = gr.CheckboxGroup(
label=category["name"],
choices=category["tags"]
)
tag_components.append(tag_component)
other_description = gr.Textbox(
label=config["free_text_fields"][0]["name"],
lines=3
)
return quality_label, *tag_components, other_description
def load_config_to_ui(config):
return (
config["quality_scale"]["name"],
config["quality_scale"]["description"],
[[item["value"], item["label"]] for item in config["quality_scale"]["scale"]],
[[cat["name"], cat["type"], ", ".join(cat["tags"])] for cat in config["tag_categories"]],
[[field["name"], field["description"]] for field in config["free_text_fields"]]
)
def save_config_from_ui(name, description, scale, categories, fields, topics, all_topics_text):
if all_topics_text.visible:
topics_list = [topic.strip() for topic in all_topics_text.split("\n") if topic.strip()]
else:
topics_list = [topic[0] for topic in topics]
new_config = {
"quality_scale": {
"name": name,
"description": description,
"scale": [{"value": row[0], "label": row[1]} for row in scale]
},
"tag_categories": [{"name": row[0], "type": row[1], "tags": row[2].split(", ")} for row in categories],
"free_text_fields": [{"name": row[0], "description": row[1]} for row in fields],
"topics": topics_list
}
save_annotation_config(new_config)
return "Configuration saved successfully", new_config
# Add this new function to generate the preview
def generate_preview(row_data, quality, high_quality_tags, low_quality_tags, toxic_tags, other):
try:
data = json.loads(row_data)
metadata = {
"annotation": {
"quality": quality,
"tags": {
"high_quality": high_quality_tags,
"low_quality": low_quality_tags,
"toxic": toxic_tags
},
"free_text": {
"Additional Notes": other
}
}
}
data["metadata"] = metadata
return json.dumps(data, indent=2)
except json.JSONDecodeError:
return "Error: Invalid JSON in the current row data"
def load_dataset_config():
params = load_params()
with open("system_messages.py", "r") as f:
system_messages_content = f.read()
vodalus_system_message = re.search(r'SYSTEM_MESSAGES_VODALUS = \[(.*?)\]', system_messages_content, re.DOTALL).group(1).strip()[3:-3]
with open("main.py", "r") as f:
main_content = f.read()
prompt_1 = re.search(r'PROMPT_1 = """(.*?)"""', main_content, re.DOTALL).group(1).strip()
topics_module = importlib.import_module("topics")
topics_list = topics_module.TOPICS
return {
"vodalus_system_message": vodalus_system_message,
"prompt_1": prompt_1,
"topics": [[topic] for topic in topics_list],
"max_tokens": params.get('max_tokens', 2048),
"temperature": params.get('temperature', 0.7),
"top_p": params.get('top_p', 0.9),
"frequency_penalty": params.get('frequency_penalty', 0.0),
"presence_penalty": params.get('presence_penalty', 0.0)
}
def edit_all_topics_func(topics):
topics_list = [topic[0] for topic in topics]
jsonl_rows = "\n".join([json.dumps({"topic": topic}) for topic in topics_list])
return (
gr.update(visible=False),
gr.update(value=jsonl_rows, visible=True),
gr.update(visible=True)
)
def update_topics_from_text(text):
try:
# Try parsing as JSONL
topics_list = [json.loads(line)["topic"] for line in text.split("\n") if line.strip()]
except json.JSONDecodeError:
# If parsing fails, treat as plain text
topics_list = [topic.strip() for topic in text.split("\n") if topic.strip()]
return gr.Dataframe.update(value=[[topic] for topic in topics_list], visible=True), gr.TextArea.update(visible=False)
def save_dataset_config(system_messages, prompt_1, topics, max_tokens, temperature, top_p, frequency_penalty, presence_penalty):
# Save VODALUS_SYSTEM_MESSAGE to system_messages.py
with open("system_messages.py", "w") as f:
f.write(f'SYSTEM_MESSAGES_VODALUS = [\n"""\n{system_messages}\n""",\n]\n')
# Save PROMPT_1 to main.py
with open("main.py", "r") as f:
main_content = f.read()
updated_main_content = re.sub(
r'PROMPT_1 = """.*?"""',
f'PROMPT_1 = """\n{prompt_1}\n"""',
main_content,
flags=re.DOTALL
)
with open("main.py", "w") as f:
f.write(updated_main_content)
# Save TOPICS to topics.py
topics_content = "TOPICS = [\n"
for topic in topics:
topics_content += f' "{topic[0]}",\n'
topics_content += "]\n"
with open("topics.py", "w") as f:
f.write(topics_content)
save_params({
'max_tokens': max_tokens,
'temperature': temperature,
'top_p': top_p,
'frequency_penalty': frequency_penalty,
'presence_penalty': presence_penalty
})
return "Dataset configuration saved successfully"
def chat_with_llm(message, history):
try:
msg_list = [{"role": "system", "content": "You are an AI assistant helping with dataset annotation and quality checking."}]
for h in history:
msg_list.append({"role": "user", "content": h[0]})
msg_list.append({"role": "assistant", "content": h[1]})
msg_list.append({"role": "user", "content": message})
# Update this line to use send_to_llm_wrapper
response, _ = send_to_llm_wrapper(msg_list)
return history + [[message, response]]
except Exception as e:
print(f"Error in chat_with_llm: {str(e)}")
return history + [[message, f"Error: {str(e)}"]]
def update_chat_context(row_data, index, total, quality, high_quality_tags, low_quality_tags, toxic_tags, other):
context = f"""Current app state:
Row: {index + 1}/{total}
Quality: {quality}
High Quality Tags: {', '.join(high_quality_tags)}
Low Quality Tags: {', '.join(low_quality_tags)}
Toxic Tags: {', '.join(toxic_tags)}
Additional Notes: {other}
Data: {row_data}
"""
return [[None, context]]
async def run_generate_dataset(num_workers, num_generations, output_file_path, llm_provider, dataset):
if loaded_dataset is None:
return "Error: No dataset loaded. Please load a dataset before generating.", ""
generated_data = []
for _ in range(num_generations):
topic_selected = random.choice(TOPICS)
system_message_selected = random.choice(SYSTEM_MESSAGES_VODALUS)
data = await generate_data(topic_selected, PROMPT_1, system_message_selected, output_file_path, llm_provider)
if data:
generated_data.append(json.dumps(data))
# Write the generated data to the output file
with open(output_file_path, 'a') as f:
for entry in generated_data:
f.write(entry + '\n')
return f"Generated {num_generations} entries and saved to {output_file_path}", "\n".join(generated_data[:5]) + "\n..."
def add_topic_row(data):
if isinstance(data, pd.DataFrame):
return pd.concat([data, pd.DataFrame({"Topic": ["New Topic"]})], ignore_index=True)
else:
return data + [["New Topic"]]
def remove_last_topic_row(data):
return data[:-1] if len(data) > 1 else data
def edit_all_topics_func(topics):
topics_list = [topic[0] for topic in topics]
jsonl_rows = "\n".join([json.dumps({"topic": topic}) for topic in topics_list])
return (
gr.update(visible=False),
gr.update(value=jsonl_rows, visible=True),
gr.update(visible=True)
)
def update_topics_from_text(text):
try:
# Try parsing as JSONL
topics_list = [json.loads(line)["topic"] for line in text.split("\n") if line.strip()]
except json.JSONDecodeError:
# If parsing fails, treat as plain text
topics_list = [topic.strip() for topic in text.split("\n") if topic.strip()]
return gr.Dataframe.update(value=[[topic] for topic in topics_list], visible=True), gr.TextArea.update(visible=False)
def update_topics_from_text(text):
try:
# Try parsing as JSONL
topics_list = [json.loads(line)["topic"] for line in text.split("\n") if line.strip()]
except json.JSONDecodeError:
# If parsing fails, treat as plain text
topics_list = [topic.strip() for topic in text.split("\n") if topic.strip()]
return gr.Dataframe.update(value=[[topic] for topic in topics_list], visible=True), gr.TextArea.update(visible=False)
def search_huggingface_datasets(query):
try:
api = HfApi()
datasets = api.list_datasets(search=query, limit=20)
dataset_ids = [dataset.id for dataset in datasets]
return gr.update(choices=dataset_ids, visible=True), ""
except Exception as e:
print(f"Error searching datasets: {str(e)}")
return gr.update(choices=["Error: Could not search datasets"], visible=True), ""
def load_huggingface_dataset(dataset_name, split="train"):
try:
print(f"Attempting to load dataset: {dataset_name}")
# Check if dataset_name is a string
if not isinstance(dataset_name, str):
raise ValueError(f"Expected dataset_name to be a string, but got {type(dataset_name)}")
# Try loading the dataset without specifying a config
full_dataset = load_dataset(dataset_name)
print(f"Dataset loaded. Available splits: {list(full_dataset.keys())}")
# Select the appropriate split
if split in full_dataset:
dataset = full_dataset[split]
print(f"Using specified split: {split}")
else:
available_splits = list(full_dataset.keys())
if available_splits:
dataset = full_dataset[available_splits[0]]
split = available_splits[0]
print(f"Specified split not found. Using first available split: {split}")
else:
raise ValueError("No valid splits found in the dataset")
return dataset, f"Dataset '{dataset_name}' (split: {split}) loaded successfully."
except Exception as e:
error_msg = f"Error loading dataset: {str(e)}"
print(f"Error details: {error_msg}")
# If loading fails, try to get the dataset card
try:
dataset_card = hf_hub_download(repo_id=dataset_name, filename="README.md")
with open(dataset_card, 'r') as f:
card_content = f.read()
return None, f"Dataset couldn't be loaded, but here's the dataset card:\n\n{card_content[:500]}..."
except:
return None, error_msg
# Wrapper function to handle the Gradio interface
def load_dataset_wrapper(dataset_name, split):
if not dataset_name:
return None, "Please enter a dataset name."
dataset, message = load_huggingface_dataset(dataset_name, split)
return dataset, message
def update_field_visibility(provider):
if provider == "local-model":
return gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
elif provider == "anything-llm":
return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)
else:
return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
def get_popular_datasets():
return [
"wikipedia",
"squad",
"glue",
"imdb",
"wmt16",
"common_voice",
"cnn_dailymail",
"amazon_reviews_multi",
"yelp_review_full",
"ag_news"
]
def load_dataset_config_for_ui():
config = load_dataset_config()
return (
config["vodalus_system_message"],
config["prompt_1"],
config["topics"],
config["max_tokens"],
config["temperature"],
config["top_p"],
config["frequency_penalty"],
config["presence_penalty"]
)
css = """
body, #root {
margin: 0;
padding: 0;
width: 100%;
height: 100%;
overflow-x: hidden;
}
.gradio-container {
max-width: 100% !important;
width: 100% !important;
margin: 0 auto !important;
padding: 0 !important;
}
.message-row {
justify-content: space-evenly !important;
}
.message-bubble-border {
border-radius: 6px !important;
}
.message-buttons-bot, .message-buttons-user {
right: 10px !important;
left: auto !important;
bottom: 2px !important;
}
.dark.message-bubble-border {
border-color: #343140 !important;
}
.dark.user {
background: #1e1c26 !important;
}
.dark.assistant.dark, .dark.pending.dark {
background: #16141c !important;
}
.tab-nav {
border-bottom: 2px solid #e0e0e0 !important;
}
.tab-nav button {
font-size: 16px !important;
padding: 10px 20px !important;
}
.input-row {
margin-bottom: 20px !important;
}
.button-row {
display: flex !important;
justify-content: space-between !important;
margin-top: 20px !important;
}
#row-editor {
height: 80vh !important;
font-size: 16px !important;
}
.file-upload-row {
height: 50px !important;
margin-bottom: 1rem !important;
}
.file-upload-row > .gr-column {
min-width: 0 !important;
}
.compact-file-upload {
height: 50px !important;
overflow: hidden !important;
}
.compact-file-upload > .file-preview {
min-height: 0 !important;
max-height: 50px !important;
padding: 0 !important;
}
.compact-file-upload > .file-preview > .file-preview-handler {
height: 50px !important;
padding: 0 8px !important;
display: flex !important;
align-items: center !important;
}
.compact-file-upload > .file-preview > .file-preview-handler > .file-preview-title {
white-space: nowrap !important;
overflow: hidden !important;
text-overflow: ellipsis !important;
flex: 1 !important;
}
.compact-file-upload > .file-preview > .file-preview-handler > .file-preview-remove {
padding: 0 !important;
min-width: 24px !important;
width: 24px !important;
height: 24px !important;
}
.compact-button {
height: 50px !important;
min-height: 40px !important;
width: 100% !important;
}
.compact-file-upload > label {
height: 50px !important;
padding: 0 8px !important;
display: flex !important;
align-items: center !important;
justify-content: left !important;
}
"""
demo = gr.Blocks(theme='Ama434/neutral-barlow', css=css)
with demo:
gr.Markdown("# Dataset Editor and Annotation Tool")
config = gr.State(load_annotation_config())
with gr.Row():
with gr.Column(min_width=1000):
with gr.Tab("Dataset Editor"):
with gr.Row(elem_classes="file-upload-row"):
with gr.Column(scale=3, min_width=400):
file_upload = gr.File(label="Upload Dataset File (.txt, .jsonl, or .csv)", elem_classes="compact-file-upload")
with gr.Column(scale=1, min_width=100):
load_button = gr.Button("Load Dataset", elem_classes="compact-button")
with gr.Row():
prev_button = gr.Button("← Previous")
row_index = gr.State(value=0)
total_rows = gr.State(value=0)
current_row_display = gr.Textbox(label="Current Row", interactive=False)
next_button = gr.Button("Next →")
with gr.Row():
with gr.Column(scale=3):
row_editor = gr.TextArea(label="Edit Row", lines=40)
with gr.Column(scale=2):
quality_label = gr.Radio(label="Relevance for Training", choices=[])
tag_components = [gr.CheckboxGroup(label=f"Tag Group {i+1}", choices=[]) for i in range(3)]
other_description = gr.Textbox(label="Additional annotations", lines=3)
# Add the AI Assistant as a dropdown
with gr.Accordion("AI Assistant", open=False):
chatbot = gr.Chatbot(height=300)
msg = gr.Textbox(label="Chat with AI Assistant")
clear = gr.Button("Clear")
with gr.Row():
to_markdown_button = gr.Button("Convert to Markdown")
to_json_button = gr.Button("Convert to JSON")
preview_button = gr.Button("Preview with Metadata")
save_row_button = gr.Button("Save Current Row", variant="primary")
preview_output = gr.TextArea(label="Preview", lines=20, interactive=False)
editor_status = gr.Textbox(label="Editor Status")
with gr.Tab("Annotation Configuration"):
with gr.Row():
with gr.Column(scale=1):
gr.Markdown("### Quality Scale")
quality_scale_name = gr.Textbox(label="Scale Name")
quality_scale_description = gr.Textbox(label="Scale Description", lines=2)
with gr.Column(scale=2):
quality_scale = gr.Dataframe(
headers=["Value", "Label"],
datatype=["str", "str"],
label="Quality Scale Options",
interactive=True,
col_count=(2, "fixed"),
row_count=(5, "dynamic"),
height=400,
wrap=True
)
gr.Markdown("### Tag Categories")
tag_categories = gr.Dataframe(
headers=["Name", "Type", "Tags"],
datatype=["str", "str", "str"],
label="Tag Categories",
interactive=True,
col_count=(3, "fixed"),
row_count=(3, "dynamic"),
height=250,
wrap=True
)
with gr.Row():
add_tag_category = gr.Button("Add Category")
remove_tag_category = gr.Button("Remove Last Category")
gr.Markdown("### Free Text Fields")
free_text_fields = gr.Dataframe(
headers=["Name", "Description"],
datatype=["str", "str"],
label="Free Text Fields",
interactive=True,
col_count=(2, "fixed"),
row_count=(2, "dynamic"),
height=300,
wrap=True
)
with gr.Row():
add_free_text_field = gr.Button("Add Field")
remove_free_text_field = gr.Button("Remove Last Field")
with gr.Row():
save_config_btn = gr.Button("Save Configuration", variant="primary")
config_status = gr.Textbox(label="Status", interactive=False)
with gr.Tab("Dataset Configuration"):
with gr.Row():
vodalus_system_message = gr.TextArea(label="System Message for JSONL Dataset", lines=10)
prompt_1 = gr.TextArea(label="Dataset Generation Prompt", lines=10)
gr.Markdown("### Topics")
with gr.Row():
with gr.Column(scale=2):
topics = gr.Dataframe(
headers=["Topic"],
datatype=["str"],
label="Topics",
interactive=True,
col_count=(1, "fixed"),
row_count=(5, "dynamic"),
height=200,
wrap=True
)
with gr.Column(scale=1):
with gr.Row():
add_topic = gr.Button("Add Topic")
remove_topic = gr.Button("Remove Last Topic")
edit_all_topics = gr.Button("Edit All Topics")
all_topics_edit = gr.TextArea(label="Edit All Topics (JSONL or Plain Text)", visible=False, lines=10)
format_info = gr.Markdown("""
Enter topics as JSONL (e.g., {"topic": "Example Topic"}) or plain text (one topic per line).
JSONL format allows for additional metadata if needed.
""", visible=False)
with gr.Row():
save_dataset_config_btn = gr.Button("Save Dataset Configuration", variant="primary")
dataset_config_status = gr.Textbox(label="Status")
# gr.Markdown("### Hugging Face Dataset")
# with gr.Row():
# dataset_search = gr.Textbox(label="Search Datasets")
# search_button = gr.Button("Search")
# dataset_input = gr.Textbox(label="Dataset Name", info="Enter a dataset name or select from search results")
# dataset_results = gr.Radio(label="Search Results", choices=[], visible=False)
# dataset_split = gr.Textbox(label="Dataset Split (optional)", value="train")
# load_dataset_button = gr.Button("Load Selected Dataset")
# dataset_status = gr.Textbox(label="Dataset Status")
# Add a state to store the loaded dataset
# loaded_dataset = gr.State(None)
with gr.Tab("Dataset Generation"):
with gr.Row():
num_workers = gr.Slider(minimum=1, maximum=10, value=1, step=1, label="Number of Workers")
num_generations = gr.Number(value=10, label="Number of Generations", precision=0)
with gr.Row():
output_file_path = gr.Textbox(label="Output File Path", value=OUTPUT_FILE_PATH)
start_generation_btn = gr.Button("Start Generation")
generation_status = gr.Textbox(label="Generation Status")
generation_output = gr.TextArea(label="Generation Output", lines=10)
with gr.Tab("LLM Configuration"):
with gr.Row():
provider = gr.Dropdown(choices=["local-model", "anything-llm"], label="LLM Provider")
with gr.Row():
base_url = gr.Textbox(label="Base URL (for local model)", visible=False)
model = gr.Textbox(label="Model (for local model)", visible=False)
with gr.Row():
workspace = gr.Textbox(label="Workspace (for AnythingLLM)", visible=False)
api_key = gr.Textbox(label="API Key (for AnythingLLM)", visible=False)
with gr.Accordion("Advanced Options", open=False):
with gr.Row():
max_tokens = gr.Slider(minimum=100, maximum=4096, value=2048, step=1, label="Max Tokens")
temperature = gr.Slider(minimum=0, maximum=1, value=0.7, step=0.01, label="Temperature")
with gr.Row():
top_p = gr.Slider(minimum=0, maximum=1, value=0.9, step=0.01, label="Top P")
frequency_penalty = gr.Slider(minimum=0, maximum=2, value=0.0, step=0.01, label="Frequency Penalty")
presence_penalty = gr.Slider(minimum=0, maximum=2, value=0.0, step=0.01, label="Presence Penalty")
save_llm_config_btn = gr.Button("Save LLM Configuration")
llm_config_status = gr.Textbox(label="LLM Config Status")
with gr.Row():
save_dataset_config_btn = gr.Button("Save Dataset Configuration", variant="primary")
dataset_config_status = gr.Textbox(label="Status")
add_topic.click(
lambda x: x + [["New Topic"]],
inputs=[topics],
outputs=[topics]
)
remove_topic.click(
lambda x: x[:-1] if len(x) > 0 else x,
inputs=[topics],
outputs=[topics]
)
edit_all_topics.click(
edit_all_topics_func,
inputs=[topics],
outputs=[topics, all_topics_edit, format_info]
)
all_topics_edit.submit(
update_topics_from_text,
inputs=[all_topics_edit],
outputs=[topics, all_topics_edit, format_info]
)
load_button.click(
load_dataset,
inputs=[file_upload],
outputs=[row_editor, row_index, total_rows, current_row_display, quality_label, *tag_components, other_description]
).then(
update_annotation_ui,
inputs=[config],
outputs=[quality_label, *tag_components, other_description]
)
prev_button.click(
navigate_rows,
inputs=[file_upload, row_index, gr.State("prev"), config],
outputs=[row_editor, row_index, total_rows, current_row_display, quality_label, *tag_components, other_description]
).then(
update_annotation_ui,
inputs=[config],
outputs=[quality_label, *tag_components, other_description]
)
next_button.click(
navigate_rows,
inputs=[file_upload, row_index, gr.State("next"), config],
outputs=[row_editor, row_index, total_rows, current_row_display, quality_label, *tag_components, other_description]
).then(
update_annotation_ui,
inputs=[config],
outputs=[quality_label, *tag_components, other_description]
)
save_row_button.click(
save_row_with_metadata,
inputs=[file_upload, row_index, row_editor, config, quality_label,
tag_components[0], tag_components[1], tag_components[2], other_description],
outputs=[editor_status]
).then(
lambda: "",
outputs=[preview_output]
)
to_markdown_button.click(
json_to_markdown,
inputs=[row_editor],
outputs=[row_editor]
)
to_json_button.click(
markdown_to_json,
inputs=[row_editor],
outputs=[row_editor]
)
demo.load(
load_config_to_ui,
inputs=[config],
outputs=[quality_scale_name, quality_scale_description, quality_scale, tag_categories, free_text_fields]
).then(
update_annotation_ui,
inputs=[config],
outputs=[quality_label, *tag_components, other_description]
)
save_config_btn.click(
save_config_from_ui,
inputs=[quality_scale_name, quality_scale_description, quality_scale, tag_categories, free_text_fields, topics, all_topics_edit],
outputs=[config_status, config]
).then(
update_annotation_ui,
inputs=[config],
outputs=[quality_label, *tag_components, other_description]
)
preview_button.click(
generate_preview,
inputs=[row_editor, quality_label, *tag_components, other_description],
outputs=[preview_output]
)
demo.load(
load_dataset_config,
outputs=[vodalus_system_message, prompt_1, topics, max_tokens, temperature, top_p, frequency_penalty, presence_penalty]
)
save_dataset_config_btn.click(
save_dataset_config,
inputs=[vodalus_system_message, prompt_1, topics, max_tokens, temperature, top_p, frequency_penalty, presence_penalty],
outputs=[dataset_config_status]
)
start_generation_btn.click(
run_generate_dataset,
inputs=[num_workers, num_generations, output_file_path, llm_provider_state],
outputs=[generation_status, generation_output]
)
demo.load(
load_llm_config,
outputs=[provider, base_url, model, workspace, api_key, max_tokens, temperature, top_p, frequency_penalty, presence_penalty]
)
save_llm_config_btn.click(
save_llm_config,
inputs=[provider, base_url, model, workspace, api_key, max_tokens, temperature, top_p, frequency_penalty, presence_penalty],
outputs=[llm_config_status]
)
msg.submit(chat_with_llm, [msg, chatbot], [chatbot])
clear.click(lambda: None, None, chatbot, queue=False)
for button in [load_button, prev_button, next_button]:
button.click(
update_chat_context,
inputs=[row_editor, row_index, total_rows, quality_label, *tag_components, other_description],
outputs=[chatbot]
)
provider.change(
lambda x: x,
inputs=[provider],
outputs=[llm_provider_state]
)
# search_button.click(
# search_huggingface_datasets,
# inputs=[dataset_search],
# outputs=[dataset_results, dataset_input]
# )
# dataset_results.change(
# lambda choice: choice,
# inputs=[dataset_results],
# outputs=[dataset_input]
# )
# load_dataset_button.click(
# load_dataset_wrapper,
# inputs=[dataset_input, dataset_split],
# outputs=[loaded_dataset, dataset_status]
# )
# Modify the start_generation_btn.click to include the loaded dataset
start_generation_btn.click(
run_generate_dataset,
inputs=[num_workers, num_generations, output_file_path, llm_provider_state],
outputs=[generation_status, generation_output]
)
demo.load(
load_dataset_config_for_ui,
outputs=[
vodalus_system_message,
prompt_1,
topics,
max_tokens,
temperature,
top_p,
frequency_penalty,
presence_penalty
]
)
if __name__ == "__main__":
demo.launch(share=True)