{ "cells": [ { "cell_type": "code", "execution_count": null, "id": "4e366e3c", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": 5, "id": "7e0c20cf", "metadata": {}, "outputs": [], "source": [ "import json\n", "import time\n", "import requests\n", "import logging\n", "from datetime import datetime\n", "# Try to import tqdm for a nice progress bar, fallback to a simple print if not installed\n", "try:\n", " from tqdm.auto import tqdm\n", " HAS_TQDM = True\n", "except ImportError:\n", " HAS_TQDM = False\n", " print(\"šŸ’” Tip: run `!pip install tqdm` in a cell for a beautiful progress bar!\")" ] }, { "cell_type": "code", "execution_count": 6, "id": "1f632033", "metadata": {}, "outputs": [], "source": [ "MODEL_NAME = \"qwen3.5:9b\" \n", "OLLAMA_API_URL = \"http://localhost:11434/api/generate\"\n", "INPUT_FILE = r\"c:\\Users\\Daniel\\Desktop\\TAM-AI-Engineer-Assessment\\root_cause\\user_conv_log\\inference_test_set_100.jsonl\"\n", "OUTPUT_FILE = r\"c:\\Users\\Daniel\\Desktop\\TAM-AI-Engineer-Assessment\\root_cause\\user_conv_log\\results_qwen3.5-9b_100.jsonl\"\n", "LOG_FILE = r\"c:\\Users\\Daniel\\Desktop\\TAM-AI-Engineer-Assessment\\root_cause\\user_conv_log\\ollama_benchmark_qwen3.5-9b_100.log\"\n" ] }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [], "source": [ "def run_inference(MODEL_NAME, OLLAMA_API_URL, INPUT_FILE, OUTPUT_FILE, LOG_FILE):\n", " COOLDOWN_SECONDS = 0.5 \n", "\n", " # ==========================================\n", " # LOGGING SETUP\n", " # ==========================================\n", " logging.basicConfig(\n", " level=logging.INFO,\n", " format='%(asctime)s - %(levelname)s - %(message)s',\n", " force=True, # <--- ADD THIS LINE FOR JUPYTER NOTEBOOKS\n", " handlers=[\n", " logging.FileHandler(LOG_FILE, encoding='utf-8'),\n", " # We only log errors to console if using TQDM, to avoid breaking the progress bar\n", " logging.StreamHandler() if not HAS_TQDM else logging.NullHandler() \n", " ]\n", " )\n", " logger = logging.getLogger(__name__)\n", " print(f\"šŸš€ Starting inference benchmark with model: {MODEL_NAME}\")\n", " logger.info(f\"Starting benchmark run. Model: {MODEL_NAME}, Input: {INPUT_FILE}\")\n", " try:\n", " with open(INPUT_FILE, 'r', encoding='utf-8') as f:\n", " lines = [line for line in f if line.strip()]\n", " except FileNotFoundError:\n", " print(f\"āŒ Error: Could not find input file {INPUT_FILE}\")\n", " logger.error(f\"Input file not found: {INPUT_FILE}\")\n", " lines = []\n", " if lines:\n", " # Clear or create the output file to start fresh\n", " with open(OUTPUT_FILE, 'w', encoding='utf-8') as f:\n", " pass\n", " # Setup progress bar\n", " iterable = tqdm(lines, desc=\"Processing Queries\", unit=\"req\") if HAS_TQDM else lines\n", " \n", " for i, line in enumerate(iterable):\n", " data = json.loads(line)\n", " tag = data.get(\"tag\", \"Unknown\")\n", " risk = data.get(\"risk\", \"Unknown\")\n", " prompt = data.get(\"prompt\", \"\")\n", " \n", " logger.info(f\"[{i+1}/{len(lines)}] Starting inference for Tag: {tag} (Risk: {risk})\")\n", " \n", " # If no TQDM, print a manual detailed progress line\n", " if not HAS_TQDM:\n", " print(f\"šŸ”„ [{i+1}/{len(lines)}] Inferencing Tag: {tag} ...\", end=\"\", flush=True)\n", " \n", " payload = {\n", " \"model\": MODEL_NAME,\n", " \"prompt\": prompt,\n", " \"stream\": False,\n", " \"options\": {\n", " \"num_ctx\": 4096 # Restricted context window\n", " }\n", " }\n", " \n", " try:\n", " start_time = time.time()\n", " response = requests.post(OLLAMA_API_URL, json=payload)\n", " req_time = time.time() - start_time\n", " \n", " if response.status_code != 200:\n", " err_msg = f\"API Error: {response.status_code} - {response.text}\"\n", " logger.error(err_msg)\n", " if HAS_TQDM: tqdm.write(f\"āŒ {err_msg}\")\n", " else: print(f\" āŒ FAILED ({response.status_code})\")\n", " \n", " response.raise_for_status()\n", " result = response.json()\n", " \n", " generated_text = result.get(\"response\", \"\")\n", " \n", " # Telemetry\n", " total_duration_ns = result.get(\"total_duration\", 0) \n", " eval_count = result.get(\"eval_count\", 0)\n", " eval_duration_ns = result.get(\"eval_duration\", 0)\n", " \n", " total_time_sec = total_duration_ns / 1e9\n", " tokens_per_sec = (eval_count / (eval_duration_ns / 1e9)) if eval_duration_ns > 0 else 0.0\n", " \n", " output_obj = {\n", " \"tag\": tag,\n", " \"risk\": risk,\n", " \"model\": MODEL_NAME,\n", " \"input_prompt\": prompt,\n", " \"generated_response\": generated_text,\n", " \"stats\": {\n", " \"tokens_per_sec\": round(tokens_per_sec, 2),\n", " \"total_time\": round(total_time_sec, 2),\n", " \"eval_count\": eval_count\n", " }\n", " }\n", " \n", " # Log success metrics\n", " success_msg = f\"Done in {round(total_time_sec, 1)}s | Speed: {round(tokens_per_sec, 1)} t/s | Tokens: {eval_count}\"\n", " logger.info(f\"[{i+1}/{len(lines)}] {success_msg}\")\n", " \n", " # Update Progress Bar description with live stats\n", " if HAS_TQDM:\n", " iterable.set_postfix_str(f\"Speed: {round(tokens_per_sec, 1)} t/s\")\n", " else:\n", " print(f\" āœ… {success_msg}\")\n", " \n", " with open(OUTPUT_FILE, 'a', encoding='utf-8') as out_f:\n", " out_f.write(json.dumps(output_obj, ensure_ascii=False) + '\\n')\n", " \n", " except requests.exceptions.RequestException as e:\n", " err_msg = f\"Network Error on query {i+1}: {e}\"\n", " logger.error(err_msg)\n", " if HAS_TQDM: tqdm.write(f\"āŒ {err_msg}\")\n", " else: print(\" āŒ NETWORK ERROR\")\n", " except Exception as e:\n", " err_msg = f\"Unexpected Error on query {i+1}: {e}\"\n", " logger.error(err_msg)\n", " if HAS_TQDM: tqdm.write(f\"āŒ {err_msg}\")\n", " else: print(\" āŒ ERROR\")\n", " \n", " # VRAM Cool-down\n", " time.sleep(COOLDOWN_SECONDS)\n", " \n", " print(f\"\\nāœ… Inference completed! Results saved to {OUTPUT_FILE}\")\n", " logger.info(\"Benchmark run finished successfully.\")" ] }, { "cell_type": "code", "execution_count": null, "id": "78a70894", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "šŸš€ Starting inference benchmark with model: qwen3.5:9b\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Processing Queries: 0%| | 0/100 [00:00