Update app.py

app.py

@@ -4,6 +4,7 @@ import warnings
 from types import SimpleNamespace
 
 import pandas as pd
+import numpy as np
 import streamlit as st
 import torch
 
@@ -112,7 +113,6 @@ with st.sidebar:
         model_help = "Default model for yield prediction."
         input_max_length_default = 400
         from task_yield.train import preprocess_df
-        from task_yield.prediction import inference_fn
 
     model_name_or_path = st.selectbox(
         "Model",
@@ -311,7 +311,20 @@ if run:
 
         if task == "yield prediction":
             # Use custom inference function for yield prediction
-            prediction = inference_fn(dataloader, model, CFG)
+            prediction = []
+            total = len(dataloader)
+            progress = st.progress(0, text="Predicting yields...")
+            info_placeholder = st.empty()
+            for i, inputs in enumerate(dataloader, start=1):
+                inputs = {k: v.to(device) for k, v in inputs.items()}
+                with torch.no_grad():
+                    y_preds = model(inputs)
+                prediction.extend(y_preds.to("cpu").numpy())
+                del y_preds
+                progress.progress(i / total, text=f"Predicting yields... {i}/{total}")
+                info_placeholder.caption(f"Processed batch {i} of {total}")
+
+            prediction = np.concatenate(prediction)
             output_df = input_df.copy()
             output_df["prediction"] = prediction
             output_df["prediction"] = output_df["prediction"].clip(lower=0.0, upper=100.0)