Update Finetuning_Pipeline.py

Finetuning_Pipeline.py

@@ -1,589 +1,589 @@
-# the code is structured to automate the process of extracting text from PDF documents, 
-# generating question-answer pairs from that text, 
-# fine-tuning a language model based on the generated data, and evaluating the model's performance.
-# The use of caching, error handling, and logging improves the robustness and maintainability of the script, making it suitable for practical applications in natural language processing and machine learning.
-
-import json
-import re
-import PyPDF2
-import time
-from openai import OpenAI
-import fitz  # PyMuPDF
-import pytesseract
-from PIL import Image
-import io
-import os
-from pathlib import Path
-# We have to check which Torch version for Xformers (2.3 -> 0.0.27)
-from torch import __version__; from packaging.version import Version as V
-from unsloth import FastLanguageModel
-import torch
-import json
-from datasets import load_dataset
-from unsloth.chat_templates import get_chat_template
-from trl import SFTTrainer
-from transformers import TrainingArguments, TrainerCallback
-from unsloth import is_bfloat16_supported
-import gc
-import logging
-import sys
-import subprocess
-import requests
-# Set the GLOO_SOCKET_IFNAME environment variable
-os.environ["GLOO_SOCKET_IFNAME"] = "lo"
-
-# Configure logging
-# logging.basicConfig(level=logging.INFO)
-# input_data = json.loads(sys.argv[1])  
-
-# # Process the data (example: just print it here)
-# response = f"Received data: {input_data}"
-def wait_for_server(max_attempts=60):
-    """Wait for the vLLM server to become available."""
-    url = "http://localhost:8000/health"
-    for attempt in range(max_attempts):
-        try:
-            response = requests.get(url)
-            if response.status_code == 200:
-                logging.info("vLLM server is ready!")
-                return True
-        except requests.exceptions.RequestException as e:
-            logging.info(f"Server not ready yet: {e}. Retrying in {2**attempt} seconds...")
-            time.sleep(2**attempt)
-
-
-def log_output(pipe, log_func):
-    """Helper function to log output from a subprocess pipe."""
-    for line in iter(pipe.readline, ''):
-        log_func(line.strip())
-
-def start_vllm_server(model_name):
-    cmd = [
-        "vllm",
-        "serve",
-        "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
-        "--gpu_memory_utilization=0.80",
-        "--max_model_len=4096",
-        "--enable-chunked-prefill=False",
-        "--num_scheduler_steps=2"
-    ]
-
-    logging.info(f"Starting vLLM server with command: {' '.join(cmd)}")
-
-    # Start the server subprocess
-    server_process = subprocess.Popen(
-        cmd,
-        stdout=subprocess.PIPE,
-        stderr=subprocess.PIPE,
-        text=True,
-        bufsize=1
-    )
-
-    # # Use threads to handle stdout and stderr in real-time
-    # threading.Thread(target=log_output, args=(server_process.stdout, logging.info), daemon=True).start()
-    # threading.Thread(target=log_output, args=(server_process.stderr, logging.error), daemon=True).start()
-
-    # Wait for the server to become ready
-    if not wait_for_server():
-        server_process.terminate()
-        raise Exception("Server failed to start in time.")
-    
-    return server_process
-
-
-try:
-    # Validate that we have an argument
-    if len(sys.argv) < 2:
-        raise ValueError("No input JSON provided")
-        
-    # Parse the JSON input from command line argument
-    input_data = json.loads(sys.argv[1])
-
-    # Validate required fields
-    required_fields = ['pdf_file', 'system_prompt', 'model_name', 'max_step', 'learning_rate', 'epochs']
-    missing_fields = [field for field in required_fields if field not in input_data]
-    if missing_fields:
-        raise ValueError(f"Missing required fields: {', '.join(missing_fields)}")
-        
-    # Your existing pipeline code here
-    # Access the fields from input_data dictionary:
-    pdf_file = input_data['pdf_file']
-    system_prompt = input_data['system_prompt']
-    model_name = input_data['model_name']
-    max_step = input_data['max_step']
-    learning_rate = input_data['learning_rate']
-    epochs = input_data['epochs']
-
-    # Rest of your pipeline implementation...
-
-except json.JSONDecodeError as e:
-    logging.error(f"Invalid JSON input: {str(e)}")
-    sys.exit(1)
-except ValueError as e:
-    logging.error(str(e))
-    sys.exit(1)
-except Exception as e:
-    logging.error(f"Pipeline error: {str(e)}")
-    sys.exit(1)
-# Initialize the OpenAI Client with your RunPod API Key and Endpoint URL
-
-def get_cache_filename(pdf_path):
-    pdf_stat = os.stat(pdf_path)
-    pdf_modified_time = pdf_stat.st_mtime
-    base_name = Path(pdf_path).stem
-    cache_filename = f"{base_name}_{pdf_modified_time}.txt"
-    cache_dir = "pdf_cache"
-    return os.path.join(cache_dir, cache_filename)
-
-def save_text_cache(cache_path, text_pages):
-    os.makedirs(os.path.dirname(cache_path), exist_ok=True)
-    with open(cache_path, 'w', encoding='utf-8') as f:
-        json.dump(text_pages, f, ensure_ascii=False, indent=2)
-
-def load_text_cache(cache_path):
-    try:
-        with open(cache_path, 'r', encoding='utf-8') as f:
-            return json.load(f)
-    except (FileNotFoundError, json.JSONDecodeError):
-        return None
-
-def extract_pdf_text_by_page(pdf_path):
-    cache_path = get_cache_filename(pdf_path)
-    cached_text = load_text_cache(cache_path)
-
-    if cached_text is not None:
-        logging.info(f"Loading text from cache: {cache_path}")
-        return cached_text
-
-    logging.info("Cache not found, extracting text from PDF.")
-    pdf_text_pages = []
-    
-    try:
-        with open(pdf_path, "rb") as pdf_file:
-            pdf_reader = PyPDF2.PdfReader(pdf_file)
-            total_pages = len(pdf_reader.pages)
-
-            for page_num in range(total_pages):
-                page = pdf_reader.pages[page_num]
-                page_text = page.extract_text() or ""
-
-                if not page_text.strip():
-                    pdf_doc = fitz.open(pdf_path)
-                    pdf_page = pdf_doc.load_page(page_num)
-                    image_list = pdf_page.get_images(full=True)
-
-                    if image_list:
-                        for img in image_list:
-                            xref = img[0]
-                            base_image = pdf_doc.extract_image(xref)
-                            image_bytes = base_image["image"]
-                            image = Image.open(io.BytesIO(image_bytes))
-                            ocr_text = pytesseract.image_to_string(image)
-                            page_text += ocr_text.strip()
-
-                    pdf_doc.close()
-                pdf_text_pages.append(page_text.strip() if page_text else "")
-        save_text_cache(cache_path, pdf_text_pages)
-    except Exception as e:
-        print(f"Error extracting PDF text: {str(e)}")
-
-    return pdf_text_pages or []
-
-
-def clean_json_response(response):
-    """
-    Attempt to extract a valid JSON array from a response, even if it is incomplete or contains unterminated strings.
-    """
-    response_str = "\n".join(response) if isinstance(response, list) else response
-    response_str = response_str.replace("```json", "").replace("```", "").strip()
-
-    # Search for a JSON array pattern using regex, even if it’s incomplete
-    json_array_match = re.search(r'(\[.*?\])', response_str, re.DOTALL)
-    
-    # Attempt to parse as JSON and handle unterminated strings by retrying with a trimmed response
-    if json_array_match:
-        json_str = json_array_match.group(1)
-
-        try:
-            json_data = json.loads(json_str)
-            return json_str  # Return valid JSON if successful
-
-        except json.JSONDecodeError as e:
-            # If error is due to unterminated strings or missing brackets, try trimming and re-validating
-            print(f"JSON extraction failed: {e}")
-            # Try a progressive trim of the response
-            for i in range(len(json_str), 0, -10):
-                try:
-                    trimmed_json_str = json_str[:i] + "]"  # Ensure it ends with a closing bracket
-                    json_data = json.loads(trimmed_json_str)
-                    return trimmed_json_str
-                except json.JSONDecodeError:
-                    continue  # Keep trimming until valid or no more retries
-
-    print("JSON extraction failed: Incomplete or malformed response.")
-    return ""  # Return empty string if parsing fails
-
-def generate_qa_from_chunk(text_chunk, retries=2, max_tokens=2500):
-    """
-    Generate QA pairs based on a text chunk, with JSON validation, chunk splitting, and retry logic.
-    """
-    if not text_chunk.strip():
-        return ""  # Return empty if the chunk is blank
-
-    if len(text_chunk) > max_tokens:
-        # Recursively split chunk into smaller sections if it exceeds max token limit
-        half = len(text_chunk) // 2
-        return generate_qa_from_chunk(text_chunk[:half], retries) + \
-               generate_qa_from_chunk(text_chunk[half:], retries)
-    prompt = f"""You are an AI assistant tasked with generating informative question-answer pairs from text-based documents.
-
-INPUT CONTEXT:
-{text_chunk}
-
-TASK:
-Generate relevant question-answer pairs from the provided text. Each pair must:
-1. Be directly based on the information in the text
-2. Include a clear, specific question
-3. Provide an accurate, complete response
-4. Follow the exact JSON format specified below
-
-OUTPUT FORMAT REQUIREMENTS:
-1. Respond ONLY with a JSON array
-2. Each object must contain exactly two fields:
-- "prompt": the question
-- "response": the complete answer
-3. Include no text outside the JSON array
-4. Follow this exact structure:
-
-[
-    {{
-        "prompt": "What is the daily allowance for a Subedar on domestic travel?",
-        "response": "A Subedar is entitled to a daily allowance of Rupees 600 for domestic travel. This allowance covers meals and minor incidental expenses."
-    }},
-    {{
-        "prompt": "How much reimbursement can be claimed for travel by train for a Lieutenant Colonel?",
-        "response": "A Lieutenant Colonel is entitled to AC 1st class travel by train. The full fare for AC 1st class is reimbursed, provided the journey is undertaken for official purposes and valid tickets are submitted."
-    }},
-    {{
-        "prompt": "What is the limit for claiming hotel accommodation reimbursement for a Havildar?",
-        "response": "A Havildar can claim up to Rupees 2,500 per night for hotel accommodations during official travel, subject to submission of valid receipts and adherence to the approved lodging limits."
-    }}
-]
-
-Generate the QA pairs now, following the exact format shown above."""
-
-    attempt = 0
-    while attempt < retries:
-        try:
-            """Query the vLLM server with retries."""
-            url = "http://localhost:8000/v1/chat/completions"
-            headers = {"Content-Type": "application/json"}
-            data = {
-                "model": f"hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
-                "messages": [
-                    {"role": "system", "content": "You are a helpful assistant."},
-                    {"role": "user", "content": prompt}
-                ]
-            }
-            response = requests.post(url, headers=headers, json=data, timeout=300)
-            if response and response.choices:
-                response_text = response.choices[0].message.content.strip()
-                
-                # Use improved JSON validation with trimming for incomplete responses
-                json_str = clean_json_response(response_text)
-                if json_str:
-                    return json_str
-                else:
-                    print("JSON response incomplete, retrying with split chunks.")
-                    half = len(text_chunk) // 2
-                    return generate_qa_from_chunk(text_chunk[:half], retries) + \
-                           generate_qa_from_chunk(text_chunk[half:], retries)
-
-        except Exception as e:
-            print(f"Attempt {attempt + 1} failed: {e}")
-            attempt += 1
-            time.sleep(5)  # Delay before retrying
-
-    print("All attempts failed for this chunk.")
-    return ""
-
-
-def process_pdf_file_by_pages(pdf_path, output_json_file, pages_per_chunk=4, overlap=1):
-    pdf_text_pages = extract_pdf_text_by_page(pdf_path)
-    chunks = [' '.join(pdf_text_pages[i:i + pages_per_chunk]) 
-              for i in range(0, len(pdf_text_pages) - pages_per_chunk + 1, pages_per_chunk - overlap)]
-
-    print(f"Total chunks to process: {len(chunks)}")
-    all_qa_pairs = []
-
-    for idx, chunk in enumerate(chunks):
-        print(f"Processing chunk {idx + 1}...")
-        try:
-            qa_pairs = generate_qa_from_chunk(chunk)
-
-            if not qa_pairs:
-                print(f"No QA pairs generated for chunk {idx + 1}")
-                continue
-
-            qa_pairs_cleaned = clean_json_response(qa_pairs)
-            if not qa_pairs_cleaned:
-                print(f"Failed to clean JSON for chunk {idx + 1}")
-                continue
-
-            try:
-                qa_pairs_json = json.loads(qa_pairs_cleaned)
-                all_qa_pairs.extend(qa_pairs_json)
-                print(f"Chunk {idx + 1} processed successfully.")
-            except json.JSONDecodeError as e:
-                print(f"JSON decoding error for chunk {idx + 1}: {e}")
-                print(f"Raw cleaned response: {qa_pairs_cleaned}")
-
-        except Exception as e:
-            print(f"Error processing chunk {idx + 1}: {e}")
-
-    if all_qa_pairs:
-        with open(output_json_file, 'w', encoding='utf-8') as json_file:
-            json.dump(all_qa_pairs, json_file, ensure_ascii=False, indent=4)
-        print(f"QA pairs saved to {output_json_file}")
-    else:
-        print("No QA pairs were successfully processed.")
-
-
-# pdf_path = "/home/ubuntu/Diksha/finetuning_1/Finetuning_Complete/Dockerization/TravelEnglish 1.pdf"
-pdf_file=input_data["pdf_file"]
-
-
-output_json_file = 'output_json.json'
-process_pdf_file_by_pages(pdf_file, output_json_file)
-print(f"QA pairs saved to {output_json_file}")
-
-#Here Starts the Finetuning Process
-
-max_seq_length = 2048 # Choose any! We auto support RoPE Scaling internally!
-dtype = None # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+
-load_in_4bit = True # Use 4bit quantization to reduce memory usage. Can be False.
-
-
-model, tokenizer = FastLanguageModel.from_pretrained(
-    model_name = "unsloth/Meta-Llama-3.1-8B-bnb-4bit",
-    max_seq_length = max_seq_length,
-    dtype = dtype,
-    load_in_4bit = load_in_4bit,
-)
-
-model = FastLanguageModel.get_peft_model(
-    model,
-    r = 16, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
-    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
-                      "gate_proj", "up_proj", "down_proj",],
-    lora_alpha = 16,
-    lora_dropout = 0, # Supports any, but = 0 is optimized
-    bias = "none",    # Supports any, but = "none" is optimized
-    # [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
-    use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
-    random_state = 3407,
-    use_rslora = False,  # We support rank stabilized LoRA
-    loftq_config = None, # And LoftQ
-)
-
-
-def load_data_from_json(file_path):
-    return load_dataset('json', data_files=file_path)['train']
-
-
-tokenizer = get_chat_template(
-    tokenizer,
-    chat_template="chatml",  # Supports zephyr, chatml, mistral, llama, alpaca, vicuna, vicuna_old, unsloth
-    mapping={"role": "from", "content": "value", "user": "human", "assistant": "gpt"},  # ShareGPT style
-    map_eos_token=True,  # Maps <|im_end|> to </s> instead
-)
-
-system_prompt=input_data["system_prompt"]
-
-def formatting_prompts_func(examples):
-    prompts = examples['prompt']
-    responses = examples['response']
-
-    text = []
-
-    for prompt, response in zip(prompts, responses):
-        prompt = str(prompt).strip()
-        response = str(response).strip()
-
-        # Ensure punctuation for the prompt
-        if not prompt.endswith(('?', '.', '!', ':')):
-            prompt += '.'
-
-        # Capitalize the response and ensure punctuation
-        response = response.capitalize()
-        if not response.endswith(('.', '?', '!')):
-            response += '.'
-
-        # Create conversation in dictionary format, including the system prompt
-        convo = [
-            {'from': 'system', 'value': system_prompt},
-            {'from': 'human', 'value': prompt},
-            {'from': 'gpt', 'value': response}
-        ]
-
-        # Apply tokenizer's chat template to format each conversation
-        text.append(tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=False))
-
-    return {"text": text}  # Return the formatted text
-
-# Load your JSON file
-
-dataset = load_data_from_json(output_json_file)
-
-# Apply the formatting function to the dataset to create the 'text' column
-dataset = dataset.map(formatting_prompts_func, batched=True) # This line is crucial to add the 'text' column
-
-# Print a sample to verify the formatting
-print(dataset["text"][0])
-
-
-
-class TrackBestModelCallback(TrainerCallback):
-
-    def __init__(self, output_dir):
-        super().__init__()
-        self.best_loss = float('inf')
-        self.best_step = 0
-        self.output_dir = output_dir
-        os.makedirs(self.output_dir, exist_ok=True)
-
-
-
-    def on_log(self, args, state, control, logs=None, **kwargs):
-        train_loss = logs.get("loss")
-        if train_loss is not None and train_loss < self.best_loss:
-            self.best_loss = train_loss
-            self.best_step = state.global_step
-
-            # Save the model to disk
-            model_path = os.path.join(self.output_dir, f"best_model_step_{self.best_step}.pt")
-            torch.save(kwargs['model'].state_dict(), model_path)
-            print(f"New best model saved at {model_path} with loss: {self.best_loss}")
-
-            # Remove the previous best model if it exists
-
-            for file in os.listdir(self.output_dir):
-                if file.startswith("best_model_step_") and file != f"best_model_step_{self.best_step}.pt":
-                    os.remove(os.path.join(self.output_dir, file))
-
-    def on_train_end(self, args, state, control, **kwargs):
-        print("Training ended. Best model is saved on disk.")
-
-trainer = SFTTrainer(
-    model=model,
-    tokenizer=tokenizer,
-    train_dataset=dataset,
-    dataset_text_field="text",
-    max_seq_length=max_seq_length,
-    dataset_num_proc=4,  # Reduced to balance between speed and CPU usage
-    packing=False,  # Re-enabled packing for efficiency
-    args=TrainingArguments(
-        per_device_train_batch_size=4,  # Reduced to improve speed
-        gradient_accumulation_steps=2,  # Reduced to increase update frequency
-        warmup_steps=100,  # Reduced warmup steps
-        num_train_epochs=epochs,
-        max_steps=max_step,
-        learning_rate=learning_rate,  # Slightly reduced for stability with smaller batch size
-        fp16=not is_bfloat16_supported(),
-        bf16=is_bfloat16_supported(),
-        logging_steps=50,
-        optim="adamw_8bit",
-        weight_decay=0.01,
-        lr_scheduler_type="cosine",
-        seed=3407,
-        output_dir="outputs",
-        max_grad_norm=1.0,
-        dataloader_num_workers=4,  # Adjusted to match dataset_num_proc
-        gradient_checkpointing=True,  # Re-enabled to save memory
-        adam_beta1=0.9,
-        adam_beta2=0.999,
-        adam_epsilon=1e-8,
-        ddp_find_unused_parameters=False,
-        report_to="none",  # Disable wandb logging if you're not using it
-    ),
-
-)
-
-
-#@title Show current memory stats
-
-gpu_stats = torch.cuda.get_device_properties(0)
-start_gpu_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
-max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)
-print(f"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.")
-print(f"{start_gpu_memory} GB of memory reserved.")
-
-
-# In your training script:
-
-output_dir = "outputs"
-best_model_tracker = TrackBestModelCallback(output_dir)
-trainer.add_callback(best_model_tracker)
-
-
-
-# Start training
-
-trainer_stats = trainer.train()
-print("Training completed.")
-print(f"Best loss: {best_model_tracker.best_loss} at step {best_model_tracker.best_step}")
-print("Final training loss:", trainer_stats.training_loss)
-
-
-#@title Show final memory and time stats
-
-used_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
-used_memory_for_lora = round(used_memory - start_gpu_memory, 3)
-used_percentage = round(used_memory         /max_memory*100, 3)
-lora_percentage = round(used_memory_for_lora/max_memory*100, 3)
-print(f"{trainer_stats.metrics['train_runtime']} seconds used for training.")
-print(f"{round(trainer_stats.metrics['train_runtime']/60, 2)} minutes used for training.")
-print(f"Peak reserved memory = {used_memory} GB.")
-print(f"Peak reserved memory for training = {used_memory_for_lora} GB.")
-print(f"Peak reserved memory % of max memory = {used_percentage} %.")
-print(f"Peak reserved memory for training % of max memory = {lora_percentage} %.")
-
-# Saving the model for VLLM
-# Function to save and push model in different formats
-def save_and_push_model(format, save_method, model_name):
-    try:
-        model.save_pretrained_merged("model", tokenizer, save_method=save_method)
-        model.push_to_hub_merged(f"PharynxAI/{model_name}", tokenizer, save_method=save_method, private=True, token=os.getenv('HF_Token'))
-        print(f"Successfully saved and pushed model in {format} format.")
-    except Exception as e:
-        print(f"Error while saving or pushing model in {format} format: {e}")
-if __name__ == "__main__":
-
-    server_process = None
-    
-    try:
-        # # Start vLLM server
-        server_process = start_vllm_server("hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4")
-
-    except Exception as e:
-        logging.error(f"An error occurred: {e}")
-        sys.exit(1)
-        
-    finally:
-        # Cleanup: terminate the server process if it exists
-        if server_process:
-            logging.info("Shutting down vLLM server...")
-            server_process.terminate()
-            try:
-                server_process.wait(timeout=5)
-            except subprocess.TimeoutExpired:
-                logging.warning("Server didn't terminate gracefully, forcing kill...")
-                server_process.kill()
-                server_process.wait()
-            logging.info("Server shutdown complete")
-
-    # Assuming input_data is defined and contains 'model_name' 
-    save_and_push_model("16bit", "merged_16bit", input_data['model_name'])
-    save_and_push_model("4bit", "merged_4bit", input_data['model_name'])
-    save_and_push_model("LoRA adapters", "lora", input_data['model_name'])
-
+# the code is structured to automate the process of extracting text from PDF documents, 
+# generating question-answer pairs from that text, 
+# fine-tuning a language model based on the generated data, and evaluating the model's performance.
+# The use of caching, error handling, and logging improves the robustness and maintainability of the script, making it suitable for practical applications in natural language processing and machine learning.
+
+import json
+import re
+import PyPDF2
+import time
+from openai import OpenAI
+import fitz  # PyMuPDF
+import pytesseract
+from PIL import Image
+import io
+import os
+from pathlib import Path
+# We have to check which Torch version for Xformers (2.3 -> 0.0.27)
+from torch import __version__; from packaging.version import Version as V
+from unsloth import FastLanguageModel
+import torch
+import json
+from datasets import load_dataset
+from unsloth.chat_templates import get_chat_template
+from trl import SFTTrainer
+from transformers import TrainingArguments, TrainerCallback
+from unsloth import is_bfloat16_supported
+import gc
+import logging
+import sys
+import subprocess
+import requests
+# Set the GLOO_SOCKET_IFNAME environment variable
+os.environ["GLOO_SOCKET_IFNAME"] = "lo"
+
+# Configure logging
+# logging.basicConfig(level=logging.INFO)
+# input_data = json.loads(sys.argv[1])  
+
+# # Process the data (example: just print it here)
+# response = f"Received data: {input_data}"
+def wait_for_server(max_attempts=60):
+    """Wait for the vLLM server to become available."""
+    url = "http://localhost:8000/health"
+    for attempt in range(max_attempts):
+        try:
+            response = requests.get(url)
+            if response.status_code == 200:
+                logging.info("vLLM server is ready!")
+                return True
+        except requests.exceptions.RequestException as e:
+            logging.info(f"Server not ready yet: {e}. Retrying in {2**attempt} seconds...")
+            time.sleep(2**attempt)
+
+
+def log_output(pipe, log_func):
+    """Helper function to log output from a subprocess pipe."""
+    for line in iter(pipe.readline, ''):
+        log_func(line.strip())
+
+def start_vllm_server(model_name):
+    cmd = [
+        "vllm",
+        "serve",
+        "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+        "--gpu_memory_utilization=0.98",
+        "--max_model_len=4096",
+        "--enable-chunked-prefill=False",
+        "--num_scheduler_steps=2"
+    ]
+
+    logging.info(f"Starting vLLM server with command: {' '.join(cmd)}")
+
+    # Start the server subprocess
+    server_process = subprocess.Popen(
+        cmd,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+        bufsize=1
+    )
+
+    # # Use threads to handle stdout and stderr in real-time
+    # threading.Thread(target=log_output, args=(server_process.stdout, logging.info), daemon=True).start()
+    # threading.Thread(target=log_output, args=(server_process.stderr, logging.error), daemon=True).start()
+
+    # Wait for the server to become ready
+    if not wait_for_server():
+        server_process.terminate()
+        raise Exception("Server failed to start in time.")
+    
+    return server_process
+
+
+try:
+    # Validate that we have an argument
+    if len(sys.argv) < 2:
+        raise ValueError("No input JSON provided")
+        
+    # Parse the JSON input from command line argument
+    input_data = json.loads(sys.argv[1])
+
+    # Validate required fields
+    required_fields = ['pdf_file', 'system_prompt', 'model_name', 'max_step', 'learning_rate', 'epochs']
+    missing_fields = [field for field in required_fields if field not in input_data]
+    if missing_fields:
+        raise ValueError(f"Missing required fields: {', '.join(missing_fields)}")
+        
+    # Your existing pipeline code here
+    # Access the fields from input_data dictionary:
+    pdf_file = input_data['pdf_file']
+    system_prompt = input_data['system_prompt']
+    model_name = input_data['model_name']
+    max_step = input_data['max_step']
+    learning_rate = input_data['learning_rate']
+    epochs = input_data['epochs']
+
+    # Rest of your pipeline implementation...
+
+except json.JSONDecodeError as e:
+    logging.error(f"Invalid JSON input: {str(e)}")
+    sys.exit(1)
+except ValueError as e:
+    logging.error(str(e))
+    sys.exit(1)
+except Exception as e:
+    logging.error(f"Pipeline error: {str(e)}")
+    sys.exit(1)
+# Initialize the OpenAI Client with your RunPod API Key and Endpoint URL
+
+def get_cache_filename(pdf_path):
+    pdf_stat = os.stat(pdf_path)
+    pdf_modified_time = pdf_stat.st_mtime
+    base_name = Path(pdf_path).stem
+    cache_filename = f"{base_name}_{pdf_modified_time}.txt"
+    cache_dir = "pdf_cache"
+    return os.path.join(cache_dir, cache_filename)
+
+def save_text_cache(cache_path, text_pages):
+    os.makedirs(os.path.dirname(cache_path), exist_ok=True)
+    with open(cache_path, 'w', encoding='utf-8') as f:
+        json.dump(text_pages, f, ensure_ascii=False, indent=2)
+
+def load_text_cache(cache_path):
+    try:
+        with open(cache_path, 'r', encoding='utf-8') as f:
+            return json.load(f)
+    except (FileNotFoundError, json.JSONDecodeError):
+        return None
+
+def extract_pdf_text_by_page(pdf_path):
+    cache_path = get_cache_filename(pdf_path)
+    cached_text = load_text_cache(cache_path)
+
+    if cached_text is not None:
+        logging.info(f"Loading text from cache: {cache_path}")
+        return cached_text
+
+    logging.info("Cache not found, extracting text from PDF.")
+    pdf_text_pages = []
+    
+    try:
+        with open(pdf_path, "rb") as pdf_file:
+            pdf_reader = PyPDF2.PdfReader(pdf_file)
+            total_pages = len(pdf_reader.pages)
+
+            for page_num in range(total_pages):
+                page = pdf_reader.pages[page_num]
+                page_text = page.extract_text() or ""
+
+                if not page_text.strip():
+                    pdf_doc = fitz.open(pdf_path)
+                    pdf_page = pdf_doc.load_page(page_num)
+                    image_list = pdf_page.get_images(full=True)
+
+                    if image_list:
+                        for img in image_list:
+                            xref = img[0]
+                            base_image = pdf_doc.extract_image(xref)
+                            image_bytes = base_image["image"]
+                            image = Image.open(io.BytesIO(image_bytes))
+                            ocr_text = pytesseract.image_to_string(image)
+                            page_text += ocr_text.strip()
+
+                    pdf_doc.close()
+                pdf_text_pages.append(page_text.strip() if page_text else "")
+        save_text_cache(cache_path, pdf_text_pages)
+    except Exception as e:
+        print(f"Error extracting PDF text: {str(e)}")
+
+    return pdf_text_pages or []
+
+
+def clean_json_response(response):
+    """
+    Attempt to extract a valid JSON array from a response, even if it is incomplete or contains unterminated strings.
+    """
+    response_str = "\n".join(response) if isinstance(response, list) else response
+    response_str = response_str.replace("```json", "").replace("```", "").strip()
+
+    # Search for a JSON array pattern using regex, even if it’s incomplete
+    json_array_match = re.search(r'(\[.*?\])', response_str, re.DOTALL)
+    
+    # Attempt to parse as JSON and handle unterminated strings by retrying with a trimmed response
+    if json_array_match:
+        json_str = json_array_match.group(1)
+
+        try:
+            json_data = json.loads(json_str)
+            return json_str  # Return valid JSON if successful
+
+        except json.JSONDecodeError as e:
+            # If error is due to unterminated strings or missing brackets, try trimming and re-validating
+            print(f"JSON extraction failed: {e}")
+            # Try a progressive trim of the response
+            for i in range(len(json_str), 0, -10):
+                try:
+                    trimmed_json_str = json_str[:i] + "]"  # Ensure it ends with a closing bracket
+                    json_data = json.loads(trimmed_json_str)
+                    return trimmed_json_str
+                except json.JSONDecodeError:
+                    continue  # Keep trimming until valid or no more retries
+
+    print("JSON extraction failed: Incomplete or malformed response.")
+    return ""  # Return empty string if parsing fails
+
+def generate_qa_from_chunk(text_chunk, retries=2, max_tokens=2500):
+    """
+    Generate QA pairs based on a text chunk, with JSON validation, chunk splitting, and retry logic.
+    """
+    if not text_chunk.strip():
+        return ""  # Return empty if the chunk is blank
+
+    if len(text_chunk) > max_tokens:
+        # Recursively split chunk into smaller sections if it exceeds max token limit
+        half = len(text_chunk) // 2
+        return generate_qa_from_chunk(text_chunk[:half], retries) + \
+               generate_qa_from_chunk(text_chunk[half:], retries)
+    prompt = f"""You are an AI assistant tasked with generating informative question-answer pairs from text-based documents.
+
+INPUT CONTEXT:
+{text_chunk}
+
+TASK:
+Generate relevant question-answer pairs from the provided text. Each pair must:
+1. Be directly based on the information in the text
+2. Include a clear, specific question
+3. Provide an accurate, complete response
+4. Follow the exact JSON format specified below
+
+OUTPUT FORMAT REQUIREMENTS:
+1. Respond ONLY with a JSON array
+2. Each object must contain exactly two fields:
+- "prompt": the question
+- "response": the complete answer
+3. Include no text outside the JSON array
+4. Follow this exact structure:
+
+[
+    {{
+        "prompt": "What is the daily allowance for a Subedar on domestic travel?",
+        "response": "A Subedar is entitled to a daily allowance of Rupees 600 for domestic travel. This allowance covers meals and minor incidental expenses."
+    }},
+    {{
+        "prompt": "How much reimbursement can be claimed for travel by train for a Lieutenant Colonel?",
+        "response": "A Lieutenant Colonel is entitled to AC 1st class travel by train. The full fare for AC 1st class is reimbursed, provided the journey is undertaken for official purposes and valid tickets are submitted."
+    }},
+    {{
+        "prompt": "What is the limit for claiming hotel accommodation reimbursement for a Havildar?",
+        "response": "A Havildar can claim up to Rupees 2,500 per night for hotel accommodations during official travel, subject to submission of valid receipts and adherence to the approved lodging limits."
+    }}
+]
+
+Generate the QA pairs now, following the exact format shown above."""
+
+    attempt = 0
+    while attempt < retries:
+        try:
+            """Query the vLLM server with retries."""
+            url = "http://localhost:8000/v1/chat/completions"
+            headers = {"Content-Type": "application/json"}
+            data = {
+                "model": f"hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+                "messages": [
+                    {"role": "system", "content": "You are a helpful assistant."},
+                    {"role": "user", "content": prompt}
+                ]
+            }
+            response = requests.post(url, headers=headers, json=data, timeout=300)
+            if response and response.choices:
+                response_text = response.choices[0].message.content.strip()
+                
+                # Use improved JSON validation with trimming for incomplete responses
+                json_str = clean_json_response(response_text)
+                if json_str:
+                    return json_str
+                else:
+                    print("JSON response incomplete, retrying with split chunks.")
+                    half = len(text_chunk) // 2
+                    return generate_qa_from_chunk(text_chunk[:half], retries) + \
+                           generate_qa_from_chunk(text_chunk[half:], retries)
+
+        except Exception as e:
+            print(f"Attempt {attempt + 1} failed: {e}")
+            attempt += 1
+            time.sleep(5)  # Delay before retrying
+
+    print("All attempts failed for this chunk.")
+    return ""
+
+
+def process_pdf_file_by_pages(pdf_path, output_json_file, pages_per_chunk=4, overlap=1):
+    pdf_text_pages = extract_pdf_text_by_page(pdf_path)
+    chunks = [' '.join(pdf_text_pages[i:i + pages_per_chunk]) 
+              for i in range(0, len(pdf_text_pages) - pages_per_chunk + 1, pages_per_chunk - overlap)]
+
+    print(f"Total chunks to process: {len(chunks)}")
+    all_qa_pairs = []
+
+    for idx, chunk in enumerate(chunks):
+        print(f"Processing chunk {idx + 1}...")
+        try:
+            qa_pairs = generate_qa_from_chunk(chunk)
+
+            if not qa_pairs:
+                print(f"No QA pairs generated for chunk {idx + 1}")
+                continue
+
+            qa_pairs_cleaned = clean_json_response(qa_pairs)
+            if not qa_pairs_cleaned:
+                print(f"Failed to clean JSON for chunk {idx + 1}")
+                continue
+
+            try:
+                qa_pairs_json = json.loads(qa_pairs_cleaned)
+                all_qa_pairs.extend(qa_pairs_json)
+                print(f"Chunk {idx + 1} processed successfully.")
+            except json.JSONDecodeError as e:
+                print(f"JSON decoding error for chunk {idx + 1}: {e}")
+                print(f"Raw cleaned response: {qa_pairs_cleaned}")
+
+        except Exception as e:
+            print(f"Error processing chunk {idx + 1}: {e}")
+
+    if all_qa_pairs:
+        with open(output_json_file, 'w', encoding='utf-8') as json_file:
+            json.dump(all_qa_pairs, json_file, ensure_ascii=False, indent=4)
+        print(f"QA pairs saved to {output_json_file}")
+    else:
+        print("No QA pairs were successfully processed.")
+
+
+# pdf_path = "/home/ubuntu/Diksha/finetuning_1/Finetuning_Complete/Dockerization/TravelEnglish 1.pdf"
+pdf_file=input_data["pdf_file"]
+
+
+output_json_file = 'output_json.json'
+process_pdf_file_by_pages(pdf_file, output_json_file)
+print(f"QA pairs saved to {output_json_file}")
+
+#Here Starts the Finetuning Process
+
+max_seq_length = 2048 # Choose any! We auto support RoPE Scaling internally!
+dtype = None # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+
+load_in_4bit = True # Use 4bit quantization to reduce memory usage. Can be False.
+
+
+model, tokenizer = FastLanguageModel.from_pretrained(
+    model_name = "unsloth/Meta-Llama-3.1-8B-bnb-4bit",
+    max_seq_length = max_seq_length,
+    dtype = dtype,
+    load_in_4bit = load_in_4bit,
+)
+
+model = FastLanguageModel.get_peft_model(
+    model,
+    r = 16, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
+    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
+                      "gate_proj", "up_proj", "down_proj",],
+    lora_alpha = 16,
+    lora_dropout = 0, # Supports any, but = 0 is optimized
+    bias = "none",    # Supports any, but = "none" is optimized
+    # [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
+    use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
+    random_state = 3407,
+    use_rslora = False,  # We support rank stabilized LoRA
+    loftq_config = None, # And LoftQ
+)
+
+
+def load_data_from_json(file_path):
+    return load_dataset('json', data_files=file_path)['train']
+
+
+tokenizer = get_chat_template(
+    tokenizer,
+    chat_template="chatml",  # Supports zephyr, chatml, mistral, llama, alpaca, vicuna, vicuna_old, unsloth
+    mapping={"role": "from", "content": "value", "user": "human", "assistant": "gpt"},  # ShareGPT style
+    map_eos_token=True,  # Maps <|im_end|> to </s> instead
+)
+
+system_prompt=input_data["system_prompt"]
+
+def formatting_prompts_func(examples):
+    prompts = examples['prompt']
+    responses = examples['response']
+
+    text = []
+
+    for prompt, response in zip(prompts, responses):
+        prompt = str(prompt).strip()
+        response = str(response).strip()
+
+        # Ensure punctuation for the prompt
+        if not prompt.endswith(('?', '.', '!', ':')):
+            prompt += '.'
+
+        # Capitalize the response and ensure punctuation
+        response = response.capitalize()
+        if not response.endswith(('.', '?', '!')):
+            response += '.'
+
+        # Create conversation in dictionary format, including the system prompt
+        convo = [
+            {'from': 'system', 'value': system_prompt},
+            {'from': 'human', 'value': prompt},
+            {'from': 'gpt', 'value': response}
+        ]
+
+        # Apply tokenizer's chat template to format each conversation
+        text.append(tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=False))
+
+    return {"text": text}  # Return the formatted text
+
+# Load your JSON file
+
+dataset = load_data_from_json(output_json_file)
+
+# Apply the formatting function to the dataset to create the 'text' column
+dataset = dataset.map(formatting_prompts_func, batched=True) # This line is crucial to add the 'text' column
+
+# Print a sample to verify the formatting
+print(dataset["text"][0])
+
+
+
+class TrackBestModelCallback(TrainerCallback):
+
+    def __init__(self, output_dir):
+        super().__init__()
+        self.best_loss = float('inf')
+        self.best_step = 0
+        self.output_dir = output_dir
+        os.makedirs(self.output_dir, exist_ok=True)
+
+
+
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        train_loss = logs.get("loss")
+        if train_loss is not None and train_loss < self.best_loss:
+            self.best_loss = train_loss
+            self.best_step = state.global_step
+
+            # Save the model to disk
+            model_path = os.path.join(self.output_dir, f"best_model_step_{self.best_step}.pt")
+            torch.save(kwargs['model'].state_dict(), model_path)
+            print(f"New best model saved at {model_path} with loss: {self.best_loss}")
+
+            # Remove the previous best model if it exists
+
+            for file in os.listdir(self.output_dir):
+                if file.startswith("best_model_step_") and file != f"best_model_step_{self.best_step}.pt":
+                    os.remove(os.path.join(self.output_dir, file))
+
+    def on_train_end(self, args, state, control, **kwargs):
+        print("Training ended. Best model is saved on disk.")
+
+trainer = SFTTrainer(
+    model=model,
+    tokenizer=tokenizer,
+    train_dataset=dataset,
+    dataset_text_field="text",
+    max_seq_length=max_seq_length,
+    dataset_num_proc=4,  # Reduced to balance between speed and CPU usage
+    packing=False,  # Re-enabled packing for efficiency
+    args=TrainingArguments(
+        per_device_train_batch_size=4,  # Reduced to improve speed
+        gradient_accumulation_steps=2,  # Reduced to increase update frequency
+        warmup_steps=100,  # Reduced warmup steps
+        num_train_epochs=epochs,
+        max_steps=max_step,
+        learning_rate=learning_rate,  # Slightly reduced for stability with smaller batch size
+        fp16=not is_bfloat16_supported(),
+        bf16=is_bfloat16_supported(),
+        logging_steps=50,
+        optim="adamw_8bit",
+        weight_decay=0.01,
+        lr_scheduler_type="cosine",
+        seed=3407,
+        output_dir="outputs",
+        max_grad_norm=1.0,
+        dataloader_num_workers=4,  # Adjusted to match dataset_num_proc
+        gradient_checkpointing=True,  # Re-enabled to save memory
+        adam_beta1=0.9,
+        adam_beta2=0.999,
+        adam_epsilon=1e-8,
+        ddp_find_unused_parameters=False,
+        report_to="none",  # Disable wandb logging if you're not using it
+    ),
+
+)
+
+
+#@title Show current memory stats
+
+gpu_stats = torch.cuda.get_device_properties(0)
+start_gpu_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
+max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)
+print(f"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.")
+print(f"{start_gpu_memory} GB of memory reserved.")
+
+
+# In your training script:
+
+output_dir = "outputs"
+best_model_tracker = TrackBestModelCallback(output_dir)
+trainer.add_callback(best_model_tracker)
+
+
+
+# Start training
+
+trainer_stats = trainer.train()
+print("Training completed.")
+print(f"Best loss: {best_model_tracker.best_loss} at step {best_model_tracker.best_step}")
+print("Final training loss:", trainer_stats.training_loss)
+
+
+#@title Show final memory and time stats
+
+used_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
+used_memory_for_lora = round(used_memory - start_gpu_memory, 3)
+used_percentage = round(used_memory         /max_memory*100, 3)
+lora_percentage = round(used_memory_for_lora/max_memory*100, 3)
+print(f"{trainer_stats.metrics['train_runtime']} seconds used for training.")
+print(f"{round(trainer_stats.metrics['train_runtime']/60, 2)} minutes used for training.")
+print(f"Peak reserved memory = {used_memory} GB.")
+print(f"Peak reserved memory for training = {used_memory_for_lora} GB.")
+print(f"Peak reserved memory % of max memory = {used_percentage} %.")
+print(f"Peak reserved memory for training % of max memory = {lora_percentage} %.")
+
+# Saving the model for VLLM
+# Function to save and push model in different formats
+def save_and_push_model(format, save_method, model_name):
+    try:
+        model.save_pretrained_merged("model", tokenizer, save_method=save_method)
+        model.push_to_hub_merged(f"PharynxAI/{model_name}", tokenizer, save_method=save_method, private=True, token=os.getenv('HF_Token'))
+        print(f"Successfully saved and pushed model in {format} format.")
+    except Exception as e:
+        print(f"Error while saving or pushing model in {format} format: {e}")
+if __name__ == "__main__":
+
+    server_process = None
+    
+    try:
+        # # Start vLLM server
+        server_process = start_vllm_server("hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4")
+
+    except Exception as e:
+        logging.error(f"An error occurred: {e}")
+        sys.exit(1)
+        
+    finally:
+        # Cleanup: terminate the server process if it exists
+        if server_process:
+            logging.info("Shutting down vLLM server...")
+            server_process.terminate()
+            try:
+                server_process.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                logging.warning("Server didn't terminate gracefully, forcing kill...")
+                server_process.kill()
+                server_process.wait()
+            logging.info("Server shutdown complete")
+
+    # Assuming input_data is defined and contains 'model_name' 
+    save_and_push_model("16bit", "merged_16bit", input_data['model_name'])
+    save_and_push_model("4bit", "merged_4bit", input_data['model_name'])
+    save_and_push_model("LoRA adapters", "lora", input_data['model_name'])
+