How I Built My Own ChatLLM to Explain Me Python in ChatGPT Tone

If you are a programmer or writer then there is a high probability that you have used ChatGPT. It is undoubtedly an amazing tool that can enhance our productivity. Currently, the model based on its predecessor GPT-3.5 is available for free as a research preview and I loved how it can generate code and explain it. I was particularly amazed by how it explains code well. Currently, it's free but we don’t know what will happen in the future. Also, there are many cases where we can’t directly use chatGPT especially if we are dealing with a proprietary code belonging to our organization.

Why build our own ChatLLM for Code?

As mentioned above, we can’t always put our code on ChatGPT due to security reasons. Another reason is that it will be much cheaper to use your model if you have a lot of use daily. According to some sources it takes around $700000 daily to run chatGPT and hence OpenAI has no other option but to raise the cost of the service in the future to sustain the service. Not to mention that if there is some outage on the OpenAI side some day the service at your end will suffer too.

Well, all of the above reasons seem promising to develop our model but why train any model at all when there is already a bunch of open-source models available on the internet for code? The reason I decided to do it is that I loved the tone in which the chatGPT explains the code is much better than any pre-trained model. Also, finetuned models suffer less from the hallucinations.

Creating the Data for training our Code explainer LLM

Probably this is the main section of the entire blog as data is the main recipe that we need to create our code explainer.

Since I am a Machine Learning engineer, I mostly use Python at my work, and since I have limited access to resources and time I decided to just develop the code explainer for Python and a few of its libraries that are widely used in Data Science. So let’s understand in detail how I generated the dataset for the model training.

The above figure summaries the data generation step. Now let's understand the detailed process of how I generated the data systematically.

Creating the dataset with basic Python programs

Most of the LLMs would already have some code in the pretraining data but to give the chatGPT tone I decided to start with some basic Python programs. For that, I scraped the popular learning platform w3schools for its example programs. I also took some random repositories which contained programs on basic Python. Finally, I asked chatGPT to generate explanations for the collected programs and generated a dataset from it.

At the time I thought of this project the Lamma 2 model was one of the best-performing open-source models and its pre-trained weight and code were easily available for finetuning hence I decided to use the 7b variation of the model for finetuning as its computationally cheap to experiment with the same. Since the model is a decoder-only language model the input and the completions are in the same input text and hence the data to use for finetuning the model should be in the format of completions. So I converted the data in the format below.

### Human: python code.### Assistant: Explanation of the code<EOS>

Note: Since the scrapping scripts were changed multiple times for the structure of different web pages, uploading them would not make any sense hence I didn’t do it. Also fine-tuning the model can be easily found on the internet hence I didn’t include any code for that part in the blog. The repository mentioned at the end contains the notebook for fine-tuning the model.

Finetuning The Model

I used one of the standard notebooks available over the internet to finetune the lamma2 7b model over my dataset. The model was trained with a batch size of 4. The PEFT and LORA were utilized for finetuning the model on a V100 instance with 16GB of GPU memory and 51 GB of RAM. At the same instance, the 13B model can be trained as well. The 13b model had better reasoning ability but it consumes more memory and time during inference. I trained the 7b model for around 3000 steps.

Developing the UI and the Server for the Code Explainer

Since I am not a web developer it was not in my expertise to develop any UI from scratch hence I used streamit to develop the same. Now since the generation of text from LLMs often takes time and they do generate words in an autoregressive fashion I decided to implement my server as a streaming API that allows one can request at the server at a time and get the words that are generated at the moment rather than waiting for generating the whole sequence which I would explain deeply in the next sections.

Developing UI

To save time and effort I took the basic chat UI from this awesome streamlit blog. I modified it to implement the streaming feature since I want to display words the words in a streaming fashion i.e. displaying the words as they are streamed from the backend server.

          import streamlit as st
import os
import json
import requests
import time
import uuid
from transformers import AutoTokenizer

#Streaming related code.

global word_num
word_num=0

tokenizer=AutoTokenizer.from_pretrained("./")

def get_word(process_id,base_url):
    data = {
    'process_id':process_id}

    json_data = json.dumps(data)
    headers = {'Content-Type': 'application/json'}
    response = requests.post(f"{base_url}/get_word_list", data=json_data, headers=headers,timeout=10)
    ret_list=response.json().get('word_list', [])
    return ret_list



# App title
st.set_page_config(page_title="💬 Python Explainer")


# Replicate Credentials
with st.sidebar:
    st.title('💬 Python Explainer')

    st.subheader('Model parameters')
    # temperature = st.sidebar.slider('temperature', min_value=0.01, max_value=5.0, value=0.1, step=0.01)
    # top_p = st.sidebar.slider('top_p', min_value=0.01, max_value=1.0, value=0.9, step=0.01)
    max_length = st.sidebar.slider('max_length', min_value=32, max_value=1024, value=128, step=8)
    # st.markdown('📖 Learn how to build this app in this [blog](https://blog.streamlit.io/how-to-build-a-llama-2-chatbot/)!')

# Store LLM generated responses
if "messages" not in st.session_state.keys():
    st.session_state.messages = [{"role": "assistant", "content": "Enter the Python Code you want to understand."}]

# Display or clear chat messages
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.write(message["content"])

def clear_chat_history():
    st.session_state.messages = [{"role": "assistant", "content": "Enter the Python Code you want to understand."}]
st.sidebar.button('Clear Chat History', on_click=clear_chat_history)


# User-provided prompt
if prompt := st.chat_input(disabled=False):
    prompt_sp=prompt.split('\n')
    prompt=[]
    for pr in prompt_sp:
        prompt.append(pr)
    prompt='  \n'.join(prompt)
    


    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        prompt_to_display=f"""```{prompt}```"""
        st.write(prompt_to_display)

# Generate a new response if last message is not from assistant
if st.session_state.messages[-1]["role"] != "assistant":
    with st.chat_message("assistant"):
        with st.spinner("Thinking..."):
            process_id = str(uuid.uuid4())
            #process_id="bbb"

            base_url="http://f156-34-138-115-63.ngrok.io"
            
            placeholder = st.empty()

            prompt_input=prompt
            max_length=max_length
            data = {
                'prompt': prompt_input,
                'max_length': max_length,
                'process_id':process_id}

            json_data = json.dumps(data)
            headers = {'Content-Type': 'application/json'}
            response = requests.post(f"{base_url}/start_word_generation", data=json_data, headers=headers,timeout=10)

            time.sleep(10)
            ret_list=get_word(process_id,base_url)
            #former_str=""

            while 2 not in ret_list:
                ret_list=get_word(process_id,base_url)
                
                
                
                words_tok=[]
                for word_tok in ret_list:
                    #full_response += word
                    words_tok.append(word_tok)


                
                full_response=tokenizer.decode(words_tok)
                # diff_str=full_response[len(former_str):]
                # print(f"Diff text={diff_str}")

                if '<EOS>'in full_response:
                    idx_eos=full_response.index('<EOS>')
                    full_response=full_response[:idx_eos]


                placeholder.markdown(full_response)
                #former_str=full_response
                time.sleep(10)

            

    message = {"role": "assistant", "content": full_response}
    st.session_state.messages.append(message)   

In the above, we start by importing the libraries required. The function get_word requests the backend server with a process_id to get the words generated by the model for that process at the time. Then we declare the UI which includes a streamlit slider and chat components.

From line 75 the actual streaming process takes place. The UI first creates a random process_id string that will be used by the backend API to identify the request for which it should return the generated words. It then makes a request with that ID to the backend server and gets the words generated by the server at that moment. It can be noticed that a delay of 10s has been added for the request to let the model generate some response since I was performing inference on a CPU server. Also, there is a limit on the request we can make in a minute in Ngrok which I was using for tunneling my server.

Developing the Server

Since I wanted to implement the streaming functionality I explored various options during that time I came across huggingface transformers Streamer class which supports streaming but at the time it was new and not so stable and had limited documentation. Also, I wanted the frontend server to fetch the words from the backend server at a period of 10 seconds hence because of this requirement and because the Ngrok free version had limitations for the number of requests I decided to implement everything from scratch. Also, another reason was I wanted to learn more about the inner workings of decoder-only models.

          from flask import Flask, request, jsonify
import threading
import time
from flask_ngrok import run_with_ngrok
import copy

from peft import PeftConfig, PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

from transformers import (
             AutoTokenizer,
             AutoModelForCausalLM,
             LogitsProcessorList,
             MinLengthLogitsProcessor,
            StoppingCriteriaList,
             MaxLengthCriteria,
        )
import torch





base_model_name = "NousResearch/Llama-2-7b-chat-hf" #path/to/your/model/or/name/on/hub"
adapter_model_name = "/content/"

model = AutoModelForCausalLM.from_pretrained(base_model_name)
model = PeftModel.from_pretrained(model, adapter_model_name)



logits_processor = LogitsProcessorList(
          [
                 MinLengthLogitsProcessor(10, eos_token_id=model.generation_config.eos_token_id),
            ]
        )

# model=model.to("cuda")

tokenizer = AutoTokenizer.from_pretrained('NousResearch/Llama-2-7b-chat-hf')

app = Flask(__name__)
run_with_ngrok(app)

global word_process_dict

# Create a thread-safe list to store words
word_process_dict = {}

# Function to continuously generate and append words

def start_generation(process_id,prompt,max_length=128):

  global word_process_dict

  if process_id in word_process_dict:
     return "AlreadyPresent"
  else:
    word_process_dict[process_id]=[]

    prefix=""
    inputs = tokenizer(
                prefix + prompt, padding=False, add_special_tokens=False, return_tensors="pt"
            )
    inputs["prompt_text"] = prompt

    input_ids = inputs["input_ids"]

    input_ids_c=input_ids.clone()

    pad_token_id = model.generation_config.pad_token_id
    eos_token_id = model.generation_config.eos_token_id


    count=0

    dec_input_ids=torch.tensor([[]])

    decoded_tokens=torch.tensor([[[]]])

    out_tok=[]
    past_key_val=None

    while count<=max_length:
      if count==0:
        #Prepare the input.
        input_dict={}
        input_dict["input_ids"]=input_ids_c
        input_dict["position_ids"]=None
        input_dict["past_key_values"]=past_key_val
        input_dict["use_cache"]=None
        input_dict["attention_mask"]=None

        outputs=model(**input_dict)
      else:
        input_dict={}
        input_dict["input_ids"]=dec_input_ids
        input_dict["position_ids"]=None
        input_dict["past_key_values"]=past_key_val
        input_dict["use_cache"]=None
        input_dict["attention_mask"]=None
        outputs=model(**input_dict)
      next_token_logits = outputs.logits[:, -1, :]
      next_tokens_scores = logits_processor(input_ids, next_token_logits)
      past_key_val=outputs.past_key_values

      next_tokens = torch.argmax(next_tokens_scores, dim=-1)
      if next_tokens==2:
        break
      out_tok.append(next_tokens)
      to_cat=next_tokens[:, None].to(torch.int32)
      decoded_tokens = torch.cat([decoded_tokens, to_cat[:, None]], dim=-1)
      decoded_tokens=decoded_tokens.to(torch.int32)

      dec_input_ids=next_tokens[:, None].to(torch.int32)

      # word=tokenizer.decode(dec_input_ids[0])
      word=dec_input_ids[0].item()

      if count>9:
        word_process_dict[process_id].append(word)

      if word==2:
        break


      count+=1

    word_process_dict[process_id].append(2)



@app.route('/start_word_generation', methods=['POST'])
def start_word_generation():
    global word_process_dict
    # Start the word generation process in a separate thread.
    data=request.get_json()
    process_id=data["process_id"]
    print(f"Process id in start_word_gen={process_id}")
    prompt=data['prompt']
    max_length=data['max_length']

    if process_id in word_process_dict:
      return jsonify({"message": "AlreadyPresent."})
    # else:
    #   word_process_dict[process_id]=[]


    word_generation_thread = threading.Thread(target=start_generation,args=(process_id,prompt,max_length))
    word_generation_thread.daemon = True
    word_generation_thread.start()
    return jsonify({"message": "Word generation process started"})

@app.route('/get_word_list', methods=['POST'])
def get_word_list():
    # Create a copy of the word list for safe retrieval
    data=request.get_json()
    process_id=data["process_id"]
    print(f"Process id in get_word_list={process_id}")
    word_list_copy = copy.copy(word_process_dict[process_id])
    #print(word_list_copy)
    return jsonify({"word_list": word_list_copy})

if __name__ == '__main__':
    app.run()
        

In the above backend server code, we start by importing a bunch of libraries and the fine-tuned model. We then create a dictionary word_process_dict that stores the generated words for a process i.e. its key will be the process_id and its value will be a list containing the generated tokens. The start_generation function implements the grid search algorithm with caching for speeding up decoding. It takes the process_id and the prompt as input and generates the words/tokens, and stores them in the word_process_dict. I also added certain conditions in the start_generation function like the first 9 tokens of the generated result won’t added in the results as it contains unuseful tokens like (‘ — — ’,#Assistant:”).

Then we define the routes start_word_generation and get_words_list. The start_word_generation route is used to initialize the generation for a process_id. The route checks if that process_id is present in word_process_dict if it's present it returns the message “AlreadyPresent” else it creates a thread that will keep on generating a word for that process id and appending it to the list corresponding to that process id in word_process_dict.

The route get_words_list is used to get the generated words at the current timestep for the corresponding process_id.

Some Output Samples

                                      Fig: Input & Initial portion of output.

                                        Fig. Generated explanation.

Further Improvements

* Models like Mistral 7B or Zephyr can be trained to improve the results.

* More data can be added.

* Pretraining on similar data can be performed.

* UI and server can be improved for scaling up.

Reference and Further Readings

GitHub - VikasOjha666/Code-Explainer

Contribute to VikasOjha666/Code-Explainer development by creating an account on GitHub.

github.com

How to build a Llama 2 chatbot

Experiment with this open-source LLM from Meta

blog.streamlit.io

Llama 2: Open Foundation and Fine-Tuned Chat Models

In this work, we develop and release Llama 2, a collection of pretrained and fine-tuned large language models (LLMs) ranging in scale from 7 billion to 70 billion parameters. Our fine-tuned LLMs, called Llama 2-Chat, are optimized for dialogue use cases. Our models outperform open-source chat models on most benchmarks we tested, and based on our human evaluations for helpfulness and safety, may be a suitable substitute for closed-source models. We provide a detailed description of our approach to fine-tuning and safety improvements of Llama 2-Chat in order to enable the community to build on our work and contribute to the responsible development of LLMs.

arxiv.org

In this blog, you learned you can create your chatGPT-like code explainer.  If you have any questions please feel free to drop in the comment sections.