CoT Generation using LangGraph for Medical Q&A

This project generates detailed, conversational "Chain of Thought" (CoT) responses to medical questions by simulating a dialogue between a patient and a doctor. The goal is to distill the reasoning capabilities of a large language model (like Deepseek-R1) into a smaller, more efficient model. I made it at the legendry time of launch of deepseek R1.

How it Works

The core of the project is a state machine built with langgraph. It orchestrates a conversation between two AI agents: a "patient" and a "doctor".

1. Initialization: The process starts with an initial patient query and a sample doctor's reply. The doctor's reply is used to set the tone and style for the AI-generated medical responses.
2. Conversational Flow:
   • The Patient Agent initiates the conversation and asks follow-up questions based on the doctor's responses.
   • The Doctor Agent provides medical explanations and answers, guided by the initial style prompt.
3. State Management: langgraph manages the state of the conversation, alternating turns between the patient and the doctor.
4. Output: The final output is a complete, multi-turn conversation that lays out a "chain of thought," making the reasoning process transparent and easier to understand. The conversation is also saved to conversation_log.txt.

Key Components

• CoT Generation (cot.py): This script uses a large model (GPT-4o) to generate high-quality conversational data. This data serves as the "knowledge" to be distilled.
• Question Generation (question.py): A utility script to generate a dataset of medical questions from a text corpus stored in a PostgreSQL database. This is used for preparing data for the system.
• Model Fine-tuning (nb/Llama3.1_(8B)-GRPO.ipynb): This Jupyter notebook contains an experimental setup for fine-tuning a smaller, more efficient model (Llama 3.1 8B) using the data generated by cot.py. The aim is to transfer the reasoning capabilities from the larger model to the smaller one.

Fine-tuning Strategy

The fine-tuning process uses a sophisticated reinforcement learning approach to enhance the reasoning abilities of the smaller model.

• Model: The notebook uses meta-llama/meta-Llama-3.1-8B-Instruct as the base model for fine-tuning.
• Framework: Unsloth is used for its high-performance kernels that significantly speed up the training process.
• Method: The fine-tuning is performed using Generative Representational Policy Optimization (GRPO), a reinforcement learning technique implemented via the TRL library. This method is more stable than traditional PPO for training LLMs.
• Data: For Testing The openai/gsm8k dataset, which contains grade-school math problems, is used to train the model's reasoning and problem-solving skills.
• Reward System: A custom multi-component reward system guides the model's learning process:
  • Correctness Reward: The primary reward, given if the model produces the correct answer.
  • Formatting Rewards: The model is encouraged to follow a specific XML-like format (<reasoning>...</reasoning><answer>...</answer>) through both strict and soft reward functions. This improves the model's ability to structure its output and follow instructions.
• Parameter-Efficient Fine-Tuning (PEFT): LoRA (Low-Rank Adaptation) is used to make the fine-tuning process memory-efficient by only training a small number of adapter weights.

This strategy allows the smaller model to learn complex reasoning tasks by being rewarded for both correctness and structured thinking, effectively distilling the capabilities of a larger model.

Platform

The fine-tuning and experimentation for this project were conducted on the Lightning AI platform. Lightning AI provides a powerful and streamlined environment for training and deploying AI models at scale.

• Scalable Compute: The platform offers on-demand access to high-performance GPUs, which is essential for training large language models.
• Simplified Workflow: Lightning AI simplifies the process of setting up the environment, managing dependencies, and running training jobs, allowing for a greater focus on model development.
• Integrated Studios: The use of pre-configured "Studios" provides a ready-to-use environment with the necessary tools and libraries, accelerating the development cycle.

Setup and Usage

1. Clone the repository:

   git clone https://github.com/your-username/COT-generation-using-langgraph.git
   cd COT-generation-using-langgraph

2. Install dependencies:

   pip install -r requirements.txt

3. Set up environment variables: Create a .env file in the root directory and add your OpenAI API key:

   OPENAI_API_KEY="your_openai_api_key"
   

   If you are using the question generation script, you will also need to add your database credentials.

4. Run the CoT generation:

   python cot.py

   The generated conversation will be printed to the console and saved in conversation_log.txt.

Technologies Used

• LangChain & LangGraph: For building the conversational state machine.
• OpenAI (GPT-4o): As the underlying language model for the patient and doctor agents.
• Unsloth & Hugging Face: For fine-tuning the Llama 3.1 model in the experimental notebook.
• Python: The primary programming language.
• Asyncpg: For asynchronous interaction with the PostgreSQL database in the question generation script.