import os
import json
import glob
import re
import pandas as pd
import gradio as gr
import spaces # <--- REQUIRED FOR HUGGING FACE ZEROGPU TIERS
from pypdf import PdfReader
import docx2txt
import speech_recognition as sr
from pydub import AudioSegment
from llama_index.core.node_parser import SentenceSplitter

# Directory mounts mapping perfectly to your storage volume
UPLOAD_DIR = "/data/raw_inputs"
PROCESSED_DIR = "/data/processed_jsonl"
MASTER_FILE = "/data/master_dataset.jsonl"

os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(PROCESSED_DIR, exist_ok=True)

# -------------------------------------------------------------
# CORE PIPELINE LOGIC (TRANSCRIBER, PARSER, CHUNKER)
# -------------------------------------------------------------
@spaces.GPU # <--- TELLS HUGGING FACE TO ALLOCATE GPU POWER FOR TRANSCRIBING
def transcribe_video_audio(file_path):
try:
gr.Info(f"🎬 Extracting track layers from {os.path.basename(file_path)}...")
audio = AudioSegment.from_file(file_path, format="mp4")
temp_wav = file_path + ".wav"
audio.set_channels(1).set_frame_rate(16000).export(temp_wav, format="wav")

recognizer = sr.Recognizer()
with sr.AudioFile(temp_wav) as source:
audio_data = recognizer.record(source)

gr.Info("🗣️ Processing Speech-to-Text conversion...")
extracted_text = recognizer.recognize_google(audio_data)

if os.path.exists(temp_wav):
os.remove(temp_wav)

return extracted_text.strip()
except Exception as e:
return f"[Audio Transcription Error]: {str(e)}"

def clean_text_formatting(text):
text = re.sub(r"\.([^ ])", r". \1", text)
while " " in text:
text = text.replace(" ", " ")
return text.strip()

def parse_incoming_file_to_text(file_path):
ext = os.path.splitext(file_path)[1].lower()
text = ""
if ext == ".txt":
with open(file_path, "r", encoding="utf-8") as f:
text = f.read()
elif ext == ".pdf":
reader = PdfReader(file_path)
for page in reader.pages:
t = page.extract_text()
if t: text += t + "\n"
elif ext == ".docx":
text = docx2txt.process(file_path)
elif ext in [".csv", ".xlsx"]:
df = pd.read_csv(file_path) if ext == ".csv" else pd.read_excel(file_path)
text = df.to_string(index=False)
elif ext in [".mp4", ".wav", ".mp3"]:
text = transcribe_video_audio(file_path)

return clean_text_formatting(text)

def structure_unsloth_rows(chunks, archetype):
rows = []
for idx, chunk in enumerate(chunks):
if "🎭 Persona" in archetype or "📚 Domain Expert" in archetype:
sys_msg = "You are an advanced interactive chatbot avatar."
if "🎭 Persona" in archetype:
sys_msg = "You are an immersive roleplay companion bot."
rows.append({
"conversations": [
{"from": "system", "value": sys_msg},
{"from": "human", "value": f"Context chunk {idx}: {chunk[:100]}..."},
{"from": "gpt", "value": chunk}
]
})
elif "🧮 Math Wizard" in archetype or "📈 Day Trading" in archetype:
instr = "Deconstruct structural math patterns or trading indicator calculations."
if "📈 Day Trading" in archetype:
instr = "Parse market metrics and technical data to extract signals."
rows.append({
"instruction": instr,
"input": f"Data segment context: {idx}",
"output": chunk
})
elif "💻 Code Assistant" in archetype:
rows.append({
"instruction": "Compile modular scripts based on requirements.",
"input": f"Code Requirements Segment: {idx}",
"output": chunk
})
else:
rows.append({"text": chunk})
return rows

@spaces.GPU # <--- DECORATES THE TOP LEVEL CONVERSION Pipeline FOR ZEROGPU STABILITY
def execute_dataset_builder_pipeline(files, archetype, enable_chunking, chunk_size, chunk_overlap):
if not files:
return "⚠️ Target file loading queue is empty. Please upload files."

for f in glob.glob(os.path.join(PROCESSED_DIR, "*.jsonl")):
os.remove(f)

total_files_compiled = 0
all_extracted_text_blocks = []

for file_obj in files:
raw_text = parse_incoming_file_to_text(file_obj.name)
if raw_text:
all_extracted_text_blocks.append(raw_text)
total_files_compiled += 1

if not all_extracted_text_blocks:
return "❌ Failed to extract content from assets."

combined_master_string = "\n\n--- FILE SPLIT ---\n\n".join(all_extracted_text_blocks)

if enable_chunking:
splitter = SentenceSplitter(chunk_size=int(chunk_size), chunk_overlap=int(chunk_overlap))
final_text_chunks = splitter.split_text(combined_master_string)
else:
final_text_chunks = all_extracted_text_blocks

formatted_dataset_objects = structure_unsloth_rows(final_text_chunks, archetype)

with open(MASTER_FILE, "w", encoding="utf-8") as master_f:
for obj in formatted_dataset_objects:
master_f.write(json.dumps(obj) + "\n")

return f"🔥 Conversion Complete!\n\n• Processed: {total_files_compiled}/{len(files)} files\n• Rows: {len(formatted_dataset_objects)}\n• Saved At: {MASTER_FILE}"

# -------------------------------------------------------------
# THEME TOGGLE SWITCH LOGIC ☀️/🌙
# -------------------------------------------------------------
def toggle_theme(current_theme):
if current_theme == "dark":
return gr.update(variant="light"), "light"
return gr.update(variant="dark"), "dark"

js_theme_switcher = """
function(theme) {
const documentElement = document.documentElement;
if (theme === 'dark') {
documentElement.classList.add('dark');
} else {
documentElement.classList.remove('dark');
}
return theme;
}
"""

archetype_choices = [
"🎭 Persona / Roleplay (e.g., Girlfriend, AI Companion)",
"📚 Domain Expert (e.g., History Expert, Legal Advisor)",
"🧮 Math Wizard (e.g., Algebra, Calculus solvers)",
"📈 Day Trading / Quant (e.g., XGBoost, Price Action Data)",
"💻 Code Assistant (e.g., Scripting, SQL Generation)",
"📖 Raw Knowledge Base (Continued Pre-Training)"
]

custom_theme = gr.themes.Default(
primary_hue="green",
secondary_hue="zinc",
neutral_hue="zinc"
)

# Removed theme and title parameters from constructor to prevent Gradio 6 layout deprecation warnings
with gr.Blocks() as demo:
ui_theme_state = gr.State("dark")

with gr.Row():
gr.HTML("<h1 style='flex-grow: 1; margin: 0; color: #22c55e;'>🦙 UN-SLOTH DATASET STUDIO</h1>")
theme_toggle_btn = gr.Button("🌓 Toggle Light/Dark Mode", scale=0, min_width=200)

gr.Markdown("Transform diverse media configurations into flawless JSONL files optimized for instant Unsloth training.")

with gr.Row():
with gr.Column(scale=1):
file_uploader = gr.File(file_count="multiple", label="📥 Drop Assets Here (.pdf, .txt, .docx, .mp4)")
archetype_dropdown = gr.Dropdown(choices=archetype_choices, value=archetype_choices[0], label="🤖 Choose Target AI Archetype Layout Mapping")

with gr.Accordion("✂️ Text Chunking Control Panel", open=True):
chunk_toggle = gr.Checkbox(value=True, label="Enable Smart Text Chunking Segmentation")
size_input = gr.Number(value=256, label="Chunk Token Size Limit", minimum=10, maximum=4096, step=1)
overlap_input = gr.Number(value=30, label="Overlap Token Boundary Buffer", minimum=0, maximum=1024, step=1)

run_btn = gr.Button("🚀 Run Conversion & Combine Files", variant="primary")

with gr.Column(scale=1):
log_monitor = gr.Textbox(label="🖥️ Core Engine Pipeline Logs", lines=15)

run_btn.click(
fn=execute_dataset_builder_pipeline,
inputs=[file_uploader, archetype_dropdown, chunk_toggle, size_input, overlap_input],
outputs=log_monitor
)

theme_toggle_btn.click(
fn=toggle_theme,
inputs=[ui_theme_state],
outputs=[theme_toggle_btn, ui_theme_state]
).then(
fn=None,
inputs=[ui_theme_state],
js=js_theme_switcher
)

demo.load(fn=lambda: "dark", outputs=ui_theme_state).then(fn=None, inputs=[ui_theme_state], js=js_theme_switcher)

if __name__ == "__main__":
# Theme configuration parameters passed to launch method matching Gradio 6 guidelines
demo.launch(theme=custom_theme)