You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
246 lines
9.3 KiB
Plaintext
246 lines
9.3 KiB
Plaintext
{
|
|
"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<?, ?req/s]"
|
|
]
|
|
},
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Processing Queries: 33%|███▎ | 33/100 [57:11<1:48:20, 97.03s/req, Speed: 23.9 t/s] "
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"if __name__ == \"__main__\":\n",
|
|
" run_inference(MODEL_NAME, OLLAMA_API_URL, INPUT_FILE, OUTPUT_FILE, LOG_FILE)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "cebbc129",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": []
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "0a651dbd",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": []
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": "ai-assessment",
|
|
"language": "python",
|
|
"name": "python3"
|
|
},
|
|
"language_info": {
|
|
"codemirror_mode": {
|
|
"name": "ipython",
|
|
"version": 3
|
|
},
|
|
"file_extension": ".py",
|
|
"mimetype": "text/x-python",
|
|
"name": "python",
|
|
"nbconvert_exporter": "python",
|
|
"pygments_lexer": "ipython3",
|
|
"version": "3.11.15"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 5
|
|
}
|