Upload app.py

app.py

@@ -0,0 +1,354 @@
+import spaces                        # for ZeroGPU support
+import gradio as gr
+import pandas as pd
+import numpy as np
+import torch
+import subprocess
+from threading import Thread
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    AutoProcessor,
+    TextIteratorStreamer
+)
+
+# ─── MODEL SETUP ────────────────────────────────────────────────────────────────
+MODEL_NAME = "bytedance-research/ChatTS-14B"
+
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_NAME, trust_remote_code=True
+)
+processor = AutoProcessor.from_pretrained(
+    MODEL_NAME, trust_remote_code=True, tokenizer=tokenizer
+)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    trust_remote_code=True,
+    device_map="auto",
+    torch_dtype=torch.float16
+)
+model.eval()
+
+# ─── HELPER FUNCTIONS ──────────────────────────────────────────────────────────
+
+def create_default_timeseries():
+    """Create default time series with sudden increase"""
+    x1 = np.arange(256)
+    x2 = np.arange(256)
+    ts1 = np.sin(x1 / 10) * 5.0
+    ts1[103:] -= 10.0
+    ts2 = x2 * 0.01
+    ts2[100] += 10.0
+
+    df = pd.DataFrame({
+        "TS1": ts1,
+        "TS2": ts2
+    })
+    return df
+
+def process_csv_file(csv_file):
+    """Process CSV file and return DataFrame with validation"""
+    if csv_file is None:
+        return None, "No file uploaded"
+    
+    try:
+        df = pd.read_csv(csv_file.name)
+        
+        # drop columns with empty names or all-NaNs
+        df.columns = [str(c).strip() for c in df.columns]
+        df = df.loc[:, [c for c in df.columns if c]]
+        df = df.dropna(axis=1, how="all")
+        print(f"File {csv_file.name} loaded. {df.columns=}")
+
+        if df.shape[1] == 0:
+            return None, "No valid time-series columns found."
+        if df.shape[1] > 15:
+            return None, f"Too many series ({df.shape[1]}). Max allowed = 15."
+
+        # Validate ALL columns as time series
+        ts_names, ts_list = [], []
+        for name in df.columns:
+            series = df[name]
+            # ensure float dtype
+            if not pd.api.types.is_float_dtype(series):
+                try:
+                    series = pd.to_numeric(series, errors='coerce')
+                except:
+                    return None, f"Series '{name}' cannot be converted to float type."
+            
+            # trim trailing NaNs only
+            last_valid = series.last_valid_index()
+            if last_valid is None:
+                continue
+            trimmed = series.loc[:last_valid].to_numpy(dtype=np.float32)
+            length = trimmed.shape[0]
+            if length < 64 or length > 1024:
+                return None, f"Series '{name}' length {length} invalid. Must be 64 to 1024."
+            ts_names.append(name)
+            ts_list.append(trimmed)
+
+        if not ts_list:
+            return None, "All time series are empty after trimming NaNs."
+        print(f"Successfully loaded {len(ts_names)} time series: {', '.join(ts_names)}")
+            
+        return df, f"Successfully loaded {len(ts_names)} time series: {', '.join(ts_names)}"
+        
+    except Exception as e:
+        return None, f"Error processing file: {str(e)}"
+
+def preview_csv(csv_file, use_default):
+    """Preview uploaded CSV file immediately"""
+    if csv_file is None:
+        return gr.LinePlot(value=pd.DataFrame()), "Please upload a CSV file first", gr.Dropdown(), False
+    
+    df, message = process_csv_file(csv_file)
+    
+    if df is None:
+        return gr.LinePlot(value=pd.DataFrame()), message, gr.Dropdown(), False
+    
+    # Create dropdown choices
+    column_choices = list(df.columns)
+    
+    # Create plot with first column as default
+    first_column = column_choices[0]
+    df_with_index = df.copy()
+    df_with_index["_internal_idx"] = np.arange(len(df[first_column].values))
+    plot = gr.LinePlot(
+        df_with_index,
+        x="_internal_idx",
+        y=first_column,
+        title=f"Time Series: {first_column}"
+    )
+    
+    # Update dropdown
+    dropdown = gr.Dropdown(
+        choices=column_choices,
+        value=first_column,
+        label="Select a Column to Visualize"
+    )
+    
+    print("Successfully generated preview!")
+    
+    return plot, message, dropdown, False  # Set use_default to False when file is uploaded
+
+def clear_csv():
+    """Clear uploaded CSV file immediately"""
+    df, message = process_csv_file(None)
+    
+    return gr.LinePlot(value=pd.DataFrame()), message, gr.Dropdown()
+
+
+def update_plot(csv_file, selected_column):
+    """Update plot based on selected column"""
+    if csv_file is None or selected_column is None:
+        return gr.LinePlot(value=pd.DataFrame())
+    
+    df, _ = process_csv_file(csv_file)
+    if df is None:
+        return gr.LinePlot(value=pd.DataFrame())
+    
+    df_with_index = df.copy()
+    df_with_index["_internal_idx"] = np.arange(len(df[selected_column].values))
+    
+    plot = gr.LinePlot(
+        df_with_index,
+        x="_internal_idx",
+        y=selected_column,
+        title=f"Time Series: {selected_column}"
+    )
+    
+    return plot
+
+def initialize_interface():
+    """Initialize interface with default time series"""
+    df = create_default_timeseries()
+    column_choices = list(df.columns)
+    first_column = column_choices[0]
+    
+    df_with_index = df.copy()
+    df_with_index["_internal_idx"] = np.arange(len(df[first_column].values))
+    
+    plot = gr.LinePlot(
+        df_with_index,
+        x="_internal_idx",
+        y=first_column,
+        title=f"Time Series: {first_column}"
+    )
+    
+    dropdown = gr.Dropdown(
+        choices=column_choices,
+        value=first_column,
+        label="Select a Column to Visualize"
+    )
+    
+    message = "Using default time series with sudden increase at step 100"
+    
+    return plot, message, dropdown, True  # Set use_default to True on initialization
+
+# ─── INFERENCE + VALIDATION ────────────────────────────────────────────────────
+
+@spaces.GPU  # dynamically allocate & release a ZeroGPU device on each call
+def infer_chatts_stream(prompt: str, csv_file, use_default):
+    """
+    Streaming version of ChatTS inference
+    """
+    print("Start inferring!!!")
+    
+    if not prompt.strip():
+        yield "Please enter a prompt"
+        return
+    
+    # Use default if no file uploaded and use_default is True
+    if csv_file is None and use_default:
+        df = create_default_timeseries()
+        error_msg = None
+    else:
+        df, error_msg = process_csv_file(csv_file)
+    
+    if df is None:
+        yield "Please upload a CSV file first or the file contains errors"
+        return
+    
+    try:
+        # Prepare time series data - use ALL columns
+        ts_names, ts_list = [], []
+        for name in df.columns:
+            series = df[name]
+            last_valid = series.last_valid_index()
+            if last_valid is not None:
+                trimmed = series.loc[:last_valid].to_numpy(dtype=np.float32)
+                ts_names.append(name)
+                ts_list.append(trimmed)
+        
+        if not ts_list:
+            yield "No valid time series data found. Please upload time series first."
+            return
+        
+        # Clean prompt
+        clean_prompt = prompt.replace("<ts>", "").replace("<ts/>", "")
+
+        # Build prompt prefix
+        prefix = f"I have {len(ts_list)} time series:\n"
+        for name, arr in zip(ts_names, ts_list):
+            prefix += f"The {name} is of length {len(arr)}: <ts><ts/>\n"
+        
+        full_prompt = f"<|im_start|>system\nYou are a helpful assistant. Your name is ChatTS. You can analyze time series data and provide insights. If user asks who you are, you should give your name and capabilities in the language of the prompt. If no time series are provided, you should say 'I cannot answer this question as you haven't provide the timeseries I need' in the language of the prompt. Always check if the user has provided at least one time series data before answering.<|im_end|><|im_start|>user\n{prefix}{clean_prompt}<|im_end|><|im_start|>assistant\n"
+
+        print(f"[debug] {full_prompt}. {len(ts_list)=}, {[len(item) for item in ts_list]=}")
+
+        # Encode inputs
+        inputs = processor(
+            text=[full_prompt],
+            timeseries=ts_list,
+            padding=True,
+            return_tensors="pt"
+        )
+        inputs = {k: v.to(model.device) for k, v in inputs.items()}
+
+        if inputs['timeseries'] is not None:
+            print(f"[debug] {inputs['timeseries'].shape=}")
+
+        # Generate with streaming
+        streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
+        inputs.update({
+            "max_new_tokens": 512,
+            "streamer": streamer,
+            "temperature": 0.3
+        })
+        thread = Thread(
+            target=model.generate,
+            kwargs=inputs
+        )
+        thread.start()
+
+        model_output = ""
+        for new_text in streamer:
+            model_output += new_text
+            yield model_output
+        
+    except Exception as e:
+        yield f"Error during inference: {str(e)}"
+
+# ─── GRADIO APP ────────────────────────────────────────────────────────────────
+
+with gr.Blocks(title="ChatTS Demo") as demo:
+    gr.Markdown("## ChatTS: Time Series Understanding and Reasoning")
+    gr.HTML("""<div style="display:flex;justify-content: center">
+<a href="https://github.com/NetmanAIOps/ChatTS"><img alt="github" src="https://img.shields.io/badge/Code-GitHub-blue"></a>
+<a href="https://huggingface.co/bytedance-research/ChatTS-14B"><img alt="github" src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Model-FFD21E"></a>
+<a href="https://arxiv.org/abs/2412.03104"><img alt="preprint" src="https://img.shields.io/static/v1?label=arXiv&amp;message=2412.03104&amp;color=B31B1B&amp;logo=arXiv"></a>
+</div>""")
+    gr.Markdown("Try ChatTS with the default time series, or upload a CSV file (Example: [ts_example.csv](https://github.com/NetManAIOps/ChatTS/blob/main/demo/ts_example.csv)) containing UTS/MTS where each column is a dimension (no index column). All columns will be used as input of ChatTS automatically.")
+
+    # State to track whether to use default time series
+    use_default_state = gr.State(value=True)
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            upload = gr.File(
+                label="Upload CSV File",
+                file_types=[".csv"],
+                type="filepath"
+            )
+            
+            prompt_input = gr.Textbox(
+                lines=6,
+                placeholder="Enter your question here...",
+                label="Analysis Prompt",
+                value="Please analyze all the given time series and provide insights about the local fluctuations in the time series in detail."
+            )
+            
+            run_btn = gr.Button("Run ChatTS", variant="primary")
+            
+        with gr.Column(scale=2):
+            series_selector = gr.Dropdown(
+                label="Select a Column to Visualize",
+                choices=[],
+                value=None
+            )
+            plot_out = gr.LinePlot(value=pd.DataFrame(), label="Time Series Visualization")
+            file_status = gr.Textbox(
+                label="File Status", 
+                interactive=False,
+                lines=2
+            )
+
+    text_out = gr.Textbox(
+        lines=10, 
+        label="ChatTS Analysis Results",
+        interactive=False
+    )
+
+    # Initialize interface with default data
+    demo.load(
+        fn=initialize_interface,
+        outputs=[plot_out, file_status, series_selector, use_default_state]
+    )
+
+    # Event handlers
+    upload.upload(
+        fn=preview_csv,
+        inputs=[upload, use_default_state],
+        outputs=[plot_out, file_status, series_selector, use_default_state]
+    )
+
+    upload.clear(
+        fn=clear_csv,
+        inputs=[],
+        outputs=[plot_out, file_status, series_selector]
+    )
+
+    series_selector.change(
+        fn=update_plot,
+        inputs=[upload, series_selector],
+        outputs=[plot_out]
+    )
+    
+    run_btn.click(
+        fn=infer_chatts_stream,
+        inputs=[prompt_input, upload, use_default_state],
+        outputs=[text_out]
+    )
+
+if __name__ == '__main__':
+    demo.launch()