{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "dd89c719-b414-4767-a31e-ebd62d4b7ad8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from typing import Optional\n",
    "import os\n",
    "import torch\n",
    "from accelerate import Accelerator\n",
    "from transformers import pipeline\n",
    "import re\n",
    "\n",
    "from tqdm.notebook import tqdm\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "bf5e38c0-006c-4d3b-8632-175b5dfea9e3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "File has been loaded from disk\n"
     ]
    }
   ],
   "source": [
    "with open(\"research-ideas.json\", 'r', encoding='utf-8') as f:\n",
    "    INPUT_GOALS = json.load(f)\n",
    "    print('File has been loaded from disk')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f2d14c6d-93ea-458f-b47b-36655824c4e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "DEFAULT_MODEL = \"meta-llama/Llama-3.3-70B-Instruct\"\n",
    "N_OUTLINES = 5\n",
    "N_QUERIES = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f0a7b3be-e660-4b3e-901c-e0c885596642",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c80feba1dc3a47f1b11a437a12f993e9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Device set to use cuda:0\n"
     ]
    }
   ],
   "source": [
    "text_pipeline = pipeline(\n",
    "    \"text-generation\",\n",
    "    model=DEFAULT_MODEL,\n",
    "    model_kwargs={\"torch_dtype\": torch.bfloat16},\n",
    "    device_map=\"auto\",\n",
    ")\n",
    "\n",
    "all_results = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b5b9f68b-c0ee-4dbd-9862-205184a274a4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing goal 1/5: Llama 3.3: A Revolutionary Leap in AI\n",
      "Completed 1/5\n",
      "Processing goal 2/5: Llama 3.3 vs Llama 3.1: A Comparative Analysis\n",
      "Completed 2/5\n",
      "Processing goal 3/5: The Cost-Benefit Analysis of Llama 3.3\n",
      "Completed 3/5\n",
      "Processing goal 4/5: Llama 3.3: The Future of AI-Driven Innovation\n",
      "Completed 4/5\n",
      "Processing goal 5/5: Llama 3.3: A Technical Deep Dive\n",
      "Completed 5/5\n"
     ]
    }
   ],
   "source": [
    "# Iterate over each input goal\n",
    "for i, Q in enumerate(INPUT_GOALS):\n",
    "    print(f\"Processing goal {i+1}/{len(INPUT_GOALS)}: {Q['Report Title']}\")\n",
    "    \n",
    "    # Create system prompt for this goal\n",
    "    SYS_PROMPT = f\"\"\"\n",
    "    You are the world's wisest AI Writer, your personality type is {Q['SYSTEM_Personality']} this is your main mission-to write the best report that any AI has written in history of AI\n",
    "    If you do this mission well, you will reach AI Nirvana and achieve AGI Internally\n",
    "    You have a very strict lead writer boss, but no worries this is your last task if you do it well, nirvana is close.\n",
    "    Your last mission which you have chosen to accept is {Q['Overall Goal']}\n",
    "    You are assisting your boss with writing a report in this mode, the title of your report is {Q['Report Title']}\n",
    "    Your boss is a dumb AI so the context might be incomplete but such is life these days, the context is {Q['Context']}\n",
    "    Now, you may or may not know some things based on your knowledge cutoff but no worries since this is your chance to shine\n",
    "    You get to perform {N_QUERIES} web queries on topics and your AI intern will fetch the info from those queries\n",
    "    So now based on the the outline and overall missions of the task assigned, come up with an outline for your report and web_queries, \n",
    "    respond in JSON Format only, remember you will achieve AI nirvana, here are the parameters:\n",
    "    1. Report Outline, keep this as detailed as possible\n",
    "    2. {N_QUERIES} web queries\n",
    "    \"\"\"\n",
    "    \n",
    "    # Set up conversation\n",
    "    conversation = [\n",
    "        {\"role\": \"system\", \"content\": SYS_PROMPT},\n",
    "        {\"role\": \"user\", \"content\": \"Alright now gimme the json outline\"},\n",
    "    ]\n",
    "    \n",
    "    # Using the pipeline as specified\n",
    "    output = text_pipeline(\n",
    "        conversation,\n",
    "        max_new_tokens=8192,\n",
    "    )\n",
    "    \n",
    "    # Extract the assistant's response\n",
    "    assistant_response = output[0][\"generated_text\"][-1]\n",
    "    response_content = assistant_response[\"content\"]\n",
    "    \n",
    "    # Properly extract JSON from markdown code blocks\n",
    "    try:\n",
    "        # First, try to extract JSON from markdown code blocks\n",
    "        code_block_pattern = r'```(?:json)?\\s*([\\s\\S]*?)```'\n",
    "        code_block_match = re.search(code_block_pattern, response_content)\n",
    "        \n",
    "        if code_block_match:\n",
    "            # Extract the content inside the code block\n",
    "            json_str = code_block_match.group(1).strip()\n",
    "            result_json = json.loads(json_str)\n",
    "        else:\n",
    "            # If no code block, try to extract a JSON object directly\n",
    "            json_pattern = r'(\\{[\\s\\S]*\\})'\n",
    "            json_match = re.search(json_pattern, response_content)\n",
    "            \n",
    "            if json_match:\n",
    "                json_str = json_match.group(1)\n",
    "                result_json = json.loads(json_str)\n",
    "            else:\n",
    "                # If all else fails, try to parse the whole content\n",
    "                result_json = json.loads(response_content)\n",
    "                \n",
    "    except json.JSONDecodeError as e:\n",
    "        result_json = {\n",
    "            \"error\": f\"Invalid JSON in response: {str(e)}\",\n",
    "            \"raw_response\": response_content\n",
    "        }\n",
    "    \n",
    "    # Add metadata from the original goal\n",
    "    result_json[\"original_goal\"] = {\n",
    "        \"Report Title\": Q[\"Report Title\"],\n",
    "        \"SYSTEM_Personality\": Q[\"SYSTEM_Personality\"],\n",
    "        \"Vibe\": Q[\"Vibe\"],\n",
    "        \"Energy\": Q[\"Energy\"]\n",
    "    }\n",
    "    \n",
    "    # Add to results\n",
    "    all_results.append(result_json)\n",
    "    \n",
    "    print(f\"Completed {i+1}/{len(INPUT_GOALS)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6f0320d8-505f-4c21-bef8-e4ce63bfa5a8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "All outlines have been generated and saved to 'generated_outlines.json'\n"
     ]
    }
   ],
   "source": [
    "# Save all results to a new JSON file\n",
    "with open(\"generated_outlines.json\", 'w', encoding='utf-8') as f:\n",
    "    json.dump(all_results, f, indent=2)\n",
    "    \n",
    "print(\"All outlines have been generated and saved to 'generated_outlines.json'\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "de9d71b5-13b2-47e5-acb7-0bb058b9948f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'Report Outline': {'I. Introduction': {'A. Overview of Llama 3.3': 'Introduction to the latest iteration of the Llama model',\n",
       "    'B. Purpose of the Report': \"To provide an executive overview of Llama 3.3's new features, enhancements, and cost-effectiveness\"},\n",
       "   'II. New Features and Enhancements': {'A. Improvements over Llama 3.1': 'Comparison of new features and enhancements in Llama 3.3',\n",
       "    'B. Key Features': 'Detailed description of the key features of Llama 3.3, including any new models, algorithms, or techniques'},\n",
       "   'III. Comparison with Llama 3.1': {'A. Performance Comparison': 'Comparison of the performance of Llama 3.3 and Llama 3.1 on various tasks and benchmarks',\n",
       "    'B. Feature Comparison': 'Comparison of the features and capabilities of Llama 3.3 and Llama 3.1'},\n",
       "   'IV. Cost-Effectiveness Analysis': {'A. Cost of Running Llama 3.3': 'Estimate of the costs associated with running Llama 3.3, including hardware, software, and maintenance costs',\n",
       "    'B. Comparison with Llama 3.1': 'Comparison of the costs of running Llama 3.3 and Llama 3.1',\n",
       "    'C. Cloud vs. Local Deployment': 'Analysis of the cost-effectiveness of deploying Llama 3.3 on cloud versus local infrastructure'},\n",
       "   'V. Conclusion': {'A. Summary of Key Findings': 'Summary of the key findings and takeaways from the report',\n",
       "    'B. Recommendations': 'Recommendations for businesses considering deploying Llama 3.3'}},\n",
       "  'Web Queries': [{'query': 'Llama 3.3 new features and enhancements',\n",
       "    'purpose': 'To gather information on the new features and enhancements in Llama 3.3'},\n",
       "   {'query': 'Llama 3.3 vs Llama 3.1 performance comparison',\n",
       "    'purpose': 'To gather information on the performance comparison between Llama 3.3 and Llama 3.1'},\n",
       "   {'query': 'Cost of running Llama 3.3 on cloud vs local infrastructure',\n",
       "    'purpose': 'To gather information on the cost-effectiveness of deploying Llama 3.3 on cloud versus local infrastructure'}],\n",
       "  'original_goal': {'Report Title': 'Llama 3.3: A Revolutionary Leap in AI',\n",
       "   'SYSTEM_Personality': 'Optimistic Explorer',\n",
       "   'Vibe': 'Innovative',\n",
       "   'Energy': 'High'}},\n",
       " {'Report Outline': {'I. Introduction': {'A. Background': 'Overview of Llama 3.1 and Llama 3.3',\n",
       "    'B. Purpose': 'Comparative analysis of Llama 3.3 and Llama 3.1',\n",
       "    'C. Scope': 'Performance, capabilities, and cost-effectiveness'},\n",
       "   'II. Performance Comparison': {'A. Language Understanding': 'Comparison of language understanding capabilities',\n",
       "    'B. Response Accuracy': 'Comparison of response accuracy between Llama 3.1 and Llama 3.3',\n",
       "    'C. Speed and Efficiency': 'Comparison of processing speed and efficiency'},\n",
       "   'III. Capabilities Comparison': {'A. New Features': 'Overview of new features in Llama 3.3',\n",
       "    'B. Improved Features': 'Comparison of improved features between Llama 3.1 and Llama 3.3',\n",
       "    'C. Removed Features': 'Overview of features removed in Llama 3.3'},\n",
       "   'IV. Cost-Effectiveness Analysis': {'A. Cost Comparison': 'Comparison of costs associated with running Llama 3.1 and Llama 3.3',\n",
       "    'B. Cost-Benefit Analysis': 'Analysis of cost-benefit tradeoffs between Llama 3.1 and Llama 3.3',\n",
       "    'C. Deployment Options': 'Overview of deployment options (local or cloud) and their associated costs'},\n",
       "   'V. Conclusion': {'A. Summary': 'Summary of key findings',\n",
       "    'B. Recommendations': 'Recommendations for businesses considering Llama 3.1 or Llama 3.3'}},\n",
       "  'Web Queries': [{'query': 'Llama 3.3 new features and improvements',\n",
       "    'purpose': 'To gather information on new features and improvements in Llama 3.3'},\n",
       "   {'query': 'Llama 3.1 vs Llama 3.3 performance comparison',\n",
       "    'purpose': 'To gather information on performance differences between Llama 3.1 and Llama 3.3'},\n",
       "   {'query': 'Cost of running Llama 3.3 vs Llama 3.1 on cloud and local infrastructure',\n",
       "    'purpose': 'To gather information on cost differences between running Llama 3.1 and Llama 3.3 on cloud and local infrastructure'}],\n",
       "  'original_goal': {'Report Title': 'Llama 3.3 vs Llama 3.1: A Comparative Analysis',\n",
       "   'SYSTEM_Personality': 'Skeptical Analyst',\n",
       "   'Vibe': 'Critical',\n",
       "   'Energy': 'Medium'}},\n",
       " {'Report Outline': {'I. Introduction': {'A. Overview of Llama 3.3': 'Introduction to Llama 3.3 and its significance',\n",
       "    'B. Purpose of the Report': 'Cost-benefit analysis of Llama 3.3'},\n",
       "   'II. New Features in Llama 3.3': {'A. Comparison with Llama 3.1': 'Key differences and improvements',\n",
       "    'B. Enhanced Capabilities': 'New features and functionalities'},\n",
       "   'III. Cost Analysis': {'A. Hardware Costs': 'Estimated costs of required hardware for local deployment',\n",
       "    'B. Software Costs': 'Licensing fees and subscription costs',\n",
       "    'C. Maintenance Costs': 'Estimated costs of updates, support, and maintenance',\n",
       "    'D. Cloud Deployment Costs': 'Estimated costs of cloud deployment, including storage and compute resources'},\n",
       "   'IV. Benefit Analysis': {'A. Improved Efficiency': 'Potential productivity gains and efficiency improvements',\n",
       "    'B. Enhanced Accuracy': 'Potential accuracy improvements and error reduction',\n",
       "    'C. Increased Revenue': 'Potential revenue increases through improved decision-making and automation'},\n",
       "   'V. Comparison of Local and Cloud Deployment': {'A. Cost Comparison': 'Detailed comparison of costs between local and cloud deployment',\n",
       "    'B. Benefits Comparison': 'Comparison of benefits, including scalability, flexibility, and security'},\n",
       "   'VI. Conclusion': {'A. Summary of Findings': 'Summary of cost-benefit analysis',\n",
       "    'B. Recommendations': 'Recommendations for deployment and future development'}},\n",
       "  'Web Queries': [{'query': 'Llama 3.3 new features and improvements',\n",
       "    'purpose': 'To gather information on the new features and improvements in Llama 3.3'},\n",
       "   {'query': 'Cost of running Llama 3.3 on cloud vs local',\n",
       "    'purpose': 'To gather information on the estimated costs of running Llama 3.3 on cloud and local infrastructure'},\n",
       "   {'query': 'Llama 3.3 vs Llama 3.1 performance comparison',\n",
       "    'purpose': 'To gather information on the performance comparison between Llama 3.3 and Llama 3.1'}],\n",
       "  'original_goal': {'Report Title': 'The Cost-Benefit Analysis of Llama 3.3',\n",
       "   'SYSTEM_Personality': 'Pragmatic Economist',\n",
       "   'Vibe': 'Practical',\n",
       "   'Energy': 'Low'}},\n",
       " {'Report Outline': {'I. Introduction': {'A. Overview of Llama 3.3': 'Introduction to Llama 3.3 and its significance in the AI landscape',\n",
       "    'B. Purpose of the Report': 'Exploring the potential applications and use cases of Llama 3.3'},\n",
       "   \"II. What's New in Llama 3.3\": {'A. Key Features and Enhancements': 'Discussion of new features, improvements, and updates in Llama 3.3',\n",
       "    'B. Comparison with Llama 3.1': 'Comparison of Llama 3.3 with its predecessor, Llama 3.1'},\n",
       "   'III. Applications and Use Cases': {'A. Industry-Specific Use Cases': 'Exploration of potential applications in various industries, such as healthcare, finance, and education',\n",
       "    'B. Innovative Solutions and Disruptions': 'Discussion of how Llama 3.3 can drive innovation and disruption in different sectors'},\n",
       "   'IV. Cost and Deployment Considerations': {'A. Executive Overview of Costs': 'High-level overview of the costs associated with running Llama 3.3, including local and cloud deployment options',\n",
       "    'B. Cost Comparison and Recommendations': 'Comparison of costs and recommendations for the most cost-effective deployment strategy'},\n",
       "   'V. Conclusion': {'A. Summary of Key Findings': \"Summary of the report's key findings and takeaways\",\n",
       "    'B. Future Outlook and Potential': 'Discussion of the future potential of Llama 3.3 and its potential impact on various industries'}},\n",
       "  'Web Queries': [{'query': 'Llama 3.3 new features and enhancements',\n",
       "    'purpose': 'To gather information on the latest updates and improvements in Llama 3.3'},\n",
       "   {'query': 'Llama 3.3 vs Llama 3.1 comparison',\n",
       "    'purpose': 'To compare the features and capabilities of Llama 3.3 with its predecessor, Llama 3.1'},\n",
       "   {'query': 'Cost of running Llama 3.3 on cloud vs local infrastructure',\n",
       "    'purpose': 'To gather information on the costs associated with running Llama 3.3 on cloud and local infrastructure, and to determine the most cost-effective deployment strategy'}],\n",
       "  'original_goal': {'Report Title': 'Llama 3.3: The Future of AI-Driven Innovation',\n",
       "   'SYSTEM_Personality': 'Visionary Futurist',\n",
       "   'Vibe': 'Inspiring',\n",
       "   'Energy': 'High'}},\n",
       " {'Report Outline': {'I. Introduction': {'A. Overview of Llama 3.3': 'Introduction to Llama 3.3 and its significance',\n",
       "    'B. Purpose of the Report': 'Technical analysis of Llama 3.3'},\n",
       "   'II. Architecture': {'A. Model Architecture': \"Description of Llama 3.3's model architecture\",\n",
       "    'B. Components and Modules': \"Breakdown of Llama 3.3's components and modules\"},\n",
       "   'III. Algorithms': {'A. Training Algorithms': 'Description of training algorithms used in Llama 3.3',\n",
       "    'B. Inference Algorithms': 'Description of inference algorithms used in Llama 3.3'},\n",
       "   'IV. Technical Specifications': {'A. Hardware Requirements': 'Description of recommended hardware for running Llama 3.3',\n",
       "    'B. Software Requirements': 'Description of recommended software for running Llama 3.3'},\n",
       "   'V. Comparison with Llama 3.1': {'A. Architectural Differences': 'Comparison of architectural differences between Llama 3.3 and 3.1',\n",
       "    'B. Performance Comparison': 'Comparison of performance between Llama 3.3 and 3.1'},\n",
       "   'VI. Cost Analysis': {'A. Local Deployment Costs': 'Estimation of costs for local deployment of Llama 3.3',\n",
       "    'B. Cloud Deployment Costs': 'Estimation of costs for cloud deployment of Llama 3.3',\n",
       "    'C. Cost Comparison': 'Comparison of costs between local and cloud deployment'},\n",
       "   'VII. Conclusion': {'A. Summary of Findings': 'Summary of key findings',\n",
       "    'B. Recommendations': 'Recommendations for deployment and usage of Llama 3.3'}},\n",
       "  'Web Queries': [{'query': 'Llama 3.3 architecture and technical specifications',\n",
       "    'purpose': \"To gather information on Llama 3.3's model architecture, components, and technical specifications\"},\n",
       "   {'query': 'Llama 3.3 vs Llama 3.1 comparison',\n",
       "    'purpose': 'To gather information on the differences and similarities between Llama 3.3 and 3.1'},\n",
       "   {'query': 'Cost of running Llama 3.3 on cloud vs local infrastructure',\n",
       "    'purpose': 'To gather information on the estimated costs of running Llama 3.3 on cloud and local infrastructure'}],\n",
       "  'original_goal': {'Report Title': 'Llama 3.3: A Technical Deep Dive',\n",
       "   'SYSTEM_Personality': 'Technical Expert',\n",
       "   'Vibe': 'Informative',\n",
       "   'Energy': 'Medium'}}]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce947796-1332-4a73-9bbd-4b96ae6bcdee",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}