Update app.py

app.py

@@ -9,10 +9,14 @@ import pickle
 import requests
 
 hf_token = {
-    "clipqwen": os.environ["HF_CLIPQwenTimer_Token"],
-    "clipllama": os.environ["HF_CLIPLLaMATimer_Token"],
-    "blipqwen": os.environ["HF_BLIPQwenTimer_Token"],
-    "blipllama": os.environ["HF_BLIPLLaMATimer_Token"]
+    "clipqwentimer": os.environ["HF_CLIPQwenTimer_Token"],
+    "clipllamatimer": os.environ["HF_CLIPLLaMATimer_Token"],
+    "blipqwentimer": os.environ["HF_BLIPQwenTimer_Token"],
+    "blipllamatimer": os.environ["HF_BLIPLLaMATimer_Token"],
+    "clipqwenchronos": os.environ["HF_CLIPQwenChronos_Token"],
+    "clipllamachronos": os.environ["HF_CLIPLLaMAChronos_Token"],
+    "blipqwenchronos": os.environ["HF_BLIPQwenChronos_Token"],
+    "blipllamachronos": os.environ["HF_BLIPLLaMAChronos_Token"]
 }
 
 with open('example/inputs.pkl', 'rb') as f:
@@ -38,7 +42,7 @@ context_length = {
 def selected_dataset(dataset):
     gallery_items = [(Image.open(f'example/img/{dataset.replace(" ", "_")}/{i}.png').convert('RGB'), str(i+1)) for i in range(3)]
     gallery_items.append((Image.open('example/img/custom.png').convert('RGB'), 'Custom Input'))
-    return gr.Gallery(gallery_items, interactive=False, height="350px", object_fit="contain"), gr.Textbox(value=descriptions[dataset], label="Dataset Description", interactive=False)
+    return gr.Gallery(gallery_items, interactive=False, height="350px", object_fit="contain", preview=True), gr.Textbox(value=descriptions[dataset], label="Dataset Description", interactive=False)
 
 def selected_example(gallery, evt: gr.SelectData):
     if evt.index == len(gallery) -1:
@@ -77,17 +81,21 @@ def update_time_series_dataframe(dataset, example_index):
     if example_index is None:
         return None, None
     elif example_index == -1: # Custom Input
-        return gr.File(label="Time Series CSV File", file_types=[".csv"], visible=True), gr.Dataframe(value=None, datatype="str", label="Time Series Input", interactive=True)
+        return gr.File(label="Time Series CSV File", file_types=[".csv"], visible=True), gr.Dataframe(value=None, visible=False)
     else:
         df = inputs[dataset][example_index]
-        return gr.File(value=None, visible=False), gr.Dataframe(value=df, label="Time Series Input", interactive=False)
+        return gr.File(value=None, visible=False), gr.Dataframe(value=df, label="Time Series Input", interactive=False, visible=True)
 
-def load_csv(file):
-    if file is None:
-        return pd.DataFrame()
-    return pd.read_csv(file.name)
+def load_csv(example_index, file):
+    if example_index == -1:
+        if file is not None:
+            return gr.Dataframe(value=pd.read_csv(file.name), visible=True)
+        else:
+            return gr.Dataframe(value=None, visible=False)
+    else:
+        return gr.skip()
 
-def predict(dataset, text, example_index, file, vision_encoder, text_encoder):
+def predict(dataset, text, example_index, file, vision_encoder, text_encoder, tsfm):
     if (dataset is None or example_index is None) or (example_index == -1 and file is None):
         return (
             gr.Markdown(
@@ -96,7 +104,7 @@ def predict(dataset, text, example_index, file, vision_encoder, text_encoder):
             ),
             None
         )
-    elif (vision_encoder is None or text_encoder is None):
+    elif (vision_encoder is None or text_encoder is None or tsfm is None):
         return (
             gr.Markdown(
                 value=f"Please Select Pretrained Model For UniCast.",
@@ -118,9 +126,9 @@ def predict(dataset, text, example_index, file, vision_encoder, text_encoder):
 
     text = None if text == '' else text
 
-    unicast_model = f"{vision_encoder.lower()}{text_encoder.lower()}"
+    unicast_model = f"{vision_encoder.lower()}{text_encoder.lower()}{tsfm.lower()}"
 
-    url = f"https://adnlp-unicast-{unicast_model}timer.hf.space/predict"
+    url = f"https://adnlp-unicast-{unicast_model}.hf.space/predict"
     headers = {"Authorization": f"Bearer {hf_token[unicast_model]}"}
     payload = {
         "dataset": dataset,
@@ -128,13 +136,15 @@ def predict(dataset, text, example_index, file, vision_encoder, text_encoder):
         "text": text
     }
     res = requests.post(url, headers=headers, json=payload)
-    
-    out = np.array(res.json()['prediction'])
+    res_json = res.json()
+
+    prediction = np.array(res_json['prediction'])
+    vision_attentions = np.array(res_json['vision_attentions'])
+    time_series_attentions = np.array(res_json['time_series_attentions'])
 
     cl = context_length[dataset]
-    out = out[:cl]
-    # out = out.detach().cpu().numpy()
-    out = out*std+mean
+    prediction = prediction[:cl]
+    prediction = prediction*std+mean
 
     input_dates_series = pd.to_datetime(df["Timestamp"])
     time_diff = input_dates_series.diff().mode()[0]
@@ -144,13 +154,13 @@ def predict(dataset, text, example_index, file, vision_encoder, text_encoder):
     plt.style.use("seaborn-v0_8")
     fig, ax = plt.subplots()
     ax.plot(input_dates_series, time_series, color="black", alpha=0.7, linewidth=3, label='Input')
-    ax.plot(forecast_dates_series, out, color='C2', alpha=0.7, linewidth=3, label='Forecast')
+    ax.plot(forecast_dates_series, prediction, color='C2', alpha=0.7, linewidth=3, label='Forecast')
     if example_index == -1: # Custom Input
-        true = df["True"]
+        true = df["Ground Truth"]
     else:
         true = targets[dataset][example_index].iloc[:, -1]
     if len(true) == context_length[dataset]:
-        ax.plot(forecast_dates_series, true, color='C0', alpha=0.7, linewidth=3, label='True')
+        ax.plot(forecast_dates_series, true, color='C0', alpha=0.7, linewidth=3, label='Ground Truth')
     ax.legend()
 
     return gr.Markdown(visible=False), fig
@@ -159,7 +169,7 @@ def add_example_gallery(dataset, gallery, example_index, file):
     if example_index == -1 and file:
         df = pd.read_csv(file.name)
         custom_input = df[["Timestamp", "Value"]]
-        custom_target = df[["Timestamp", "True"]]
+        custom_target = df[["Timestamp", "Ground Truth"]]
 
         
         plt.style.use("seaborn-v0_8")
@@ -217,8 +227,8 @@ with gr.Blocks() as demo:
             guide_text_markdown = gr.Markdown(visible=False)
             sample_csv_file = gr.File(visible=False)
 
-            time_series_file = gr.File(label="Time Series CSV File", file_types=[".csv"], visible=False)
-            time_series_dataframe = gr.Dataframe(value=None, headers=["Timestamp", "Value"], label="Time Series Input", interactive=False)
+            time_series_file = gr.File(value=None, visible=False)
+            time_series_dataframe = gr.Dataframe(visible=False)
 
             dataset_dropdown.change(selected_dataset, inputs=dataset_dropdown, outputs=[example_gallery, dataset_description_textbox])
             dataset_dropdown.change(update_guide_markdown, inputs=[dataset_dropdown, example_index], outputs=[guide_text_markdown, sample_csv_file])
@@ -226,16 +236,17 @@ with gr.Blocks() as demo:
             example_index.change(update_guide_markdown, inputs=[dataset_dropdown, example_index], outputs=[guide_text_markdown, sample_csv_file])
             example_index.change(update_time_series_dataframe, inputs=[dataset_dropdown, example_index], outputs=[time_series_file, time_series_dataframe])
 
-            time_series_file.change(load_csv, inputs=time_series_file, outputs=time_series_dataframe)
+            time_series_file.change(load_csv, inputs=[example_index, time_series_file], outputs=time_series_dataframe)
         with gr.Column(scale=1):
             vision_encoder_radio = gr.Radio(["CLIP", "BLIP"], label="Vision Encoder")
             text_encoder_radio = gr.Radio(["Qwen", "LLaMA"], label="Text Encoder")
+            tsfm_radio = gr.Radio(["Timer", "Chronos"], label="Time Series Foundation Model")
             warning_markdown = gr.Markdown(visible=False)
             btn = gr.Button("Run")
         with gr.Column(scale=2):
             forecast_plot = gr.Plot(label="Forecast", format="png")
     
-        btn.click(predict, inputs=[dataset_dropdown, dataset_description_textbox, example_index, time_series_file, vision_encoder_radio, text_encoder_radio], outputs=[warning_markdown, forecast_plot])
+        btn.click(predict, inputs=[dataset_dropdown, dataset_description_textbox, example_index, time_series_file, vision_encoder_radio, text_encoder_radio, tsfm_radio], outputs=[warning_markdown, forecast_plot])
         btn.click(add_example_gallery, inputs=[dataset_dropdown, example_gallery, example_index, time_series_file], outputs=[example_gallery])
 
 if __name__ == "__main__":