File size: 1,530 Bytes
975dbdd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | import os
import sys
import time
import traceback
import pandas as pd
# Ensure local package imports work
ROOT = os.path.dirname(__file__)
if ROOT not in sys.path:
sys.path.insert(0, ROOT)
import run_inference as ri
def predict_and_time(input_csv: str, output_csv: str):
df = pd.read_csv(input_csv)
results = []
for idx, row in df.iterrows():
q = row.get('question')
ground = row.get('action_label')
try:
start = time.perf_counter()
# fast routing-only prediction using loaded embedder + model
with ri.torch.no_grad():
emb = ri.embedder.encode(q, convert_to_tensor=True).to(ri.DEVICE)
action, log_prob, _ = ri.model.get_action(emb)
pred_idx = int(action.item())
end = time.perf_counter()
elapsed = end - start
except Exception:
pred_idx = None
elapsed = None
traceback.print_exc()
results.append({
'question': q,
'ground_truth': ground,
'predicted_label': pred_idx,
'time_took': elapsed,
})
out_df = pd.DataFrame(results)
out_df.to_csv(output_csv, index=False)
print(f"Wrote results to {output_csv} ({len(out_df)} rows)")
if __name__ == '__main__':
inp = os.path.join(ROOT, 'eval_queries.csv')
out = os.path.join(ROOT, 'eval_with_timing_results.csv')
print(f"Input: {inp}\nOutput: {out}\nDevice: {ri.DEVICE}")
predict_and_time(inp, out)
|