Update ChatTS

app.py

@@ -29,23 +29,115 @@ model.eval()
 
 
 # ─── INFERENCE + VALIDATION ────────────────────────────────────────────────────
-@spaces.GPU
-def generate_text(prompt):
-    inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
-    outputs = model.generate(
-        **inputs, 
-        max_new_tokens=512,
-        do_sample=True,
-        temperature=0.2,
-        top_p=0.9
+
+@spaces.GPU  # dynamically allocate & release a ZeroGPU device on each call :contentReference[oaicite:0]{index=0}
+def infer_chatts(prompt: str, csv_file):
+    """
+    1. Load CSV: first column = timestamp, other columns = TS names.
+    2. Drop empty / all-NaN columns; enforce <=15 series.
+    3. For each series: trim trailing NaNs, enforce 64 ≤ length ≤ 1024.
+    4. Build prompt prefix as:
+         I have N time series:
+         The <name> is of length L: <ts><ts/>
+    5. Encode & generate (max_new_tokens=512).
+    6. Return Gradio LinePlot + generated text.
+    """
+    # Add chat_template
+    prompt = prompt.replace("<ts>", "").replace("<ts/>", "")
+    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 time series data before answering.<|im_end|><|im_start|>user\n{prompt}<|im_end|><|im_start|>assistant\n"
+
+    # ——— CSV LOADING & PREPROCESSING ——————————————————————————————
+    df = pd.read_csv(csv_file.name, parse_dates=True, index_col=0)
+
+    # 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")
+
+    if df.shape[1] == 0:
+        raise gr.Error("No valid time-series columns found.")
+    if df.shape[1] > 15:
+        raise gr.Error(f"Too many series ({df.shape[1]}). Max allowed = 15.")
+
+    ts_names, ts_list = [], []
+    for name in df.columns:
+        series = df[name]
+        # ensure float dtype
+        if not pd.api.types.is_float_dtype(series):
+            raise gr.Error(f"Series '{name}' must be 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:
+            raise gr.Error(
+                f"Series '{name}' length {length} invalid. Must be 64 to 1024."
+            )
+        ts_names.append(name)
+        ts_list.append(trimmed)
+
+    if not ts_list:
+        raise gr.Error("All series are empty after trimming NaNs.")
+
+    # ——— BUILD PROMPT ——————————————————————————————————————————————
+    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 = prefix + prompt
+
+    # ——— ENCODE & GENERATE ————————————————————————————————————————
+    inputs = processor(
+        text=[full_prompt],
+        timeseries=ts_list,
+        padding=True,
+        return_tensors="pt"
     )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    inputs = {k: v.to(model.device) for k, v in inputs.items()}
 
-demo = gr.Interface(
-    fn=generate_text,
-    inputs=gr.Textbox(lines=2, label="Prompt"),
-    outputs=gr.Textbox(lines=6, label="Generated Text")
-)
+    outputs = model.generate(**inputs, max_new_tokens=512)
+    generated = tokenizer.decode(
+        outputs[0][ inputs["input_ids"].shape[-1] : ],
+        skip_special_tokens=True
+    )
+
+    # ——— VISUALIZATION ————————————————————————————————————————————
+    plot = gr.LinePlot(
+        df.reset_index(),
+        x=df.index.name or df.reset_index().columns[0],
+        y=ts_names,
+        label="Uploaded Time Series"
+    )
+
+    return plot, generated
+
+
+# ─── GRADIO APP ────────────────────────────────────────────────────────────────
+
+with gr.Blocks() as demo:
+    gr.Markdown("## ChatTS: Text + Time Series Inference Demo")
+
+    with gr.Row():
+        prompt_input = gr.Textbox(
+            lines=3,
+            placeholder="Enter your analysis prompt…",
+            label="Prompt"
+        )
+        upload = gr.UploadButton(
+            "Upload CSV (timestamp + float columns)",
+            file_types=[".csv"]
+        )
+
+    plot_out = gr.LinePlot(label="Time Series Visualization")
+    text_out = gr.Textbox(lines=8, label="Model Response")
+
+    run_btn = gr.Button("Run ChatTS")
+    run_btn.click(
+        fn=infer_chatts,
+        inputs=[prompt_input, upload],
+        outputs=[plot_out, text_out]
+    )
 
 
 if __name__ == '__main__':

app_chatts.py

@@ -1,141 +0,0 @@
-import spaces                        # for ZeroGPU support
-import gradio as gr
-import pandas as pd
-import numpy as np
-import torch
-import subprocess
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    AutoProcessor,
-)
-
-# ─── 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()
-
-
-# ─── INFERENCE + VALIDATION ────────────────────────────────────────────────────
-
-@spaces.GPU  # dynamically allocate & release a ZeroGPU device on each call :contentReference[oaicite:0]{index=0}
-def infer_chatts(prompt: str, csv_file):
-    """
-    1. Load CSV: first column = timestamp, other columns = TS names.
-    2. Drop empty / all-NaN columns; enforce <=15 series.
-    3. For each series: trim trailing NaNs, enforce 64 ≤ length ≤ 1024.
-    4. Build prompt prefix as:
-         I have N time series:
-         The <name> is of length L: <ts><ts/>
-    5. Encode & generate (max_new_tokens=512).
-    6. Return Gradio LinePlot + generated text.
-    """
-    # ——— CSV LOADING & PREPROCESSING ——————————————————————————————
-    df = pd.read_csv(csv_file.name, parse_dates=True, index_col=0)
-
-    # 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")
-
-    if df.shape[1] == 0:
-        raise gr.Error("No valid time-series columns found.")
-    if df.shape[1] > 15:
-        raise gr.Error(f"Too many series ({df.shape[1]}). Max allowed = 15.")
-
-    ts_names, ts_list = [], []
-    for name in df.columns:
-        series = df[name]
-        # ensure float dtype
-        if not pd.api.types.is_float_dtype(series):
-            raise gr.Error(f"Series '{name}' must be 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:
-            raise gr.Error(
-                f"Series '{name}' length {length} invalid. Must be 64 to 1024."
-            )
-        ts_names.append(name)
-        ts_list.append(trimmed)
-
-    if not ts_list:
-        raise gr.Error("All series are empty after trimming NaNs.")
-
-    # ——— BUILD PROMPT ——————————————————————————————————————————————
-    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 = prefix + prompt
-
-    # ——— ENCODE & GENERATE ————————————————————————————————————————
-    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()}
-
-    outputs = model.generate(**inputs, max_new_tokens=512)
-    generated = tokenizer.decode(
-        outputs[0][ inputs["input_ids"].shape[-1] : ],
-        skip_special_tokens=True
-    )
-
-    # ——— VISUALIZATION ————————————————————————————————————————————
-    plot = gr.LinePlot(
-        df.reset_index(),
-        x=df.index.name or df.reset_index().columns[0],
-        y=ts_names,
-        label="Uploaded Time Series"
-    )
-
-    return plot, generated
-
-
-# ─── GRADIO APP ────────────────────────────────────────────────────────────────
-
-with gr.Blocks() as demo:
-    gr.Markdown("## ChatTS: Text + Time Series Inference Demo")
-
-    with gr.Row():
-        prompt_input = gr.Textbox(
-            lines=3,
-            placeholder="Enter your analysis prompt…",
-            label="Prompt"
-        )
-        upload = gr.UploadButton(
-            "Upload CSV (timestamp + float columns)",
-            file_types=[".csv"]
-        )
-
-    plot_out = gr.LinePlot(label="Time Series Visualization")
-    text_out = gr.Textbox(lines=8, label="Model Response")
-
-    run_btn = gr.Button("Run ChatTS")
-    run_btn.click(
-        fn=infer_chatts,
-        inputs=[prompt_input, upload],
-        outputs=[plot_out, text_out]
-    )
-
-
-if __name__ == '__main__':
-    subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)
-    demo.launch()

app_legacy.py

@@ -0,0 +1,53 @@
+import spaces                        # for ZeroGPU support
+import gradio as gr
+import pandas as pd
+import numpy as np
+import torch
+import subprocess
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    AutoProcessor,
+)
+
+# ─── 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()
+
+
+# ─── INFERENCE + VALIDATION ────────────────────────────────────────────────────
+@spaces.GPU
+def generate_text(prompt):
+    inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
+    outputs = model.generate(
+        **inputs, 
+        max_new_tokens=512,
+        do_sample=True,
+        temperature=0.2,
+        top_p=0.9
+    )
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+demo = gr.Interface(
+    fn=generate_text,
+    inputs=gr.Textbox(lines=2, label="Prompt"),
+    outputs=gr.Textbox(lines=6, label="Generated Text")
+)
+
+
+if __name__ == '__main__':
+    subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)
+    demo.launch()