Utiliser une IA locale pour documenter du code

7 novembre 2023 21:47 dans data-pipelines / publications

La puissance des modèles de language open source augmente, ouvrant des possibilités nouvelles. Dans cet exemple nous allons utiliser Mistral 7B instruct, un des meilleurs modèle de taille modeste à ce jour, pour générer des docstrings pour du code Python. Un avantage essentiel de cette technique est qu'elle permet de garder le code privé, rien n'est envoyé sur des plateformes externes, assurant une confidentialité compatible avec une utilisation en entreprise.

Slider Image

Utiliser une IA locale pour documenter du code

La puissance des modèles de language open source augmente, ouvrant des possibilités nouvelles. Dans cet exemple nous allons utiliser Mistral 7B instruct, un des meilleurs modèle de taille modeste à ce jour, pour générer des docstrings pour du code Python. Un avantage essentiel de cette technique est qu'elle permet de garder le code privé, rien n'est envoyé sur des plateformes externes, assurant une confidentialité compatible avec une utilisation en entreprise.

Retrouvez le code de cet exemple dans la section exemples de la lib LocalLm.

Mise en place

Installons la librairie LocalLm:


            pip install locallm
          

Créons un dossier models et téléchargeons le modèle quantitizé (ne nécessite que 8Go de Ram et peux tourner sur Cpu):

mkdir models
            cd models
            wget https://huggingface.co/TheBloke/Mistral-7B-Instruct-v0.1-GGUF/resolve/main/mistral-7b-instruct-v0.1.Q4_K_M.gguf
          

Créons d'abord le code pour récupérer les sources du Python que l'on souhaite documenter (fonction ou classe), en utilisant inspect:

import inspect
            from importlib import import_module

            path = "mypackage.mymodule"
            node = "myfunction"
            source = inspect.getsource(getattr(import_module(path), node))
          

Chargement du modèle et paramètres

Maintenant les sources à documenter récupées, testons un premier query simple avec notre language model local. Chargeons le modèle:

from locallm import LmParams, LocalLm

            models_dir = "/absolute/path/to/my/models/dir"
            lm = LocalLm(
                LmParams(models_dir=models_dir)
            )
            lm.load_model("mistral-7b-instruct-v0.1.Q4_K_M.gguf", 8192) 
            # 8192 est la taille de la fenêtre de context du modèle (en nombre de tokens)

Note sur le deuxième paramètre de la méthode load_model: il s'agit de la taille de la fenêtre de contexte du modèle, c'est à dire le nombre de tokens total qu'il peux traiter. Le token est une unité sémantique de base du language model, qui corresponds à un mot court ou une fraction de mot (75% en moyenne). Ainsi la fenêtre représente le nombre de tokens du prompt plus le nombre de tokens générés par le modèle lors de l'inférence. Plus la fenêtre de contexte est grande, plus le modèle peux traiter de texte.

A noter l'évolution rapide vers des fenêtres de contexte plus larges: Llama 1 disposait de 2048 tokens, Llama 2 de 4096, et le modèle Mistral que nous utilisons ici dispose de 8192 tokens, éventuellement extensible. C'est largement suffisant pour traiter des portions de code significatives.

Prompt et templating

Commençons à construire un template, en utilisant le format de adapté au modèle: Mistral. A noter: le format du prompt template est essentiel et doit correspondre au format attendu par le modèle pour une efficacité optimale. Détaillons des instructions:


            template = """<s>[INST] in Python create a detailed and helpful Google style docstring for this code:

            ```python
            {prompt}
            ```
            
            Important: always provide a short example in the docstring. Break lines at 86 characters maximum. Output only the docstring [/INST]"""

La variable {prompt} sera remplacée par le code source que nous avons récupéré précedemment, construisant ainsi le prompt définitif qui sera passé au modèle. Lançons le query:


            lm.infer(
                source, # code précedemment récupéré
                InferenceParams(
                    template=template,
                    temperature=0,
                    stream=True,
                    max_tokens=1024,
                ),
            )
          

Avec le paramètre stream la librairie va afficher les tokens produits par le modèle au fur et à mesure dans le terminal. Nous utilisons une temperature à zéro pour un déterminisme maximum, étant donné la nature de la tâche.

Testons notre template avec ce code source passé en prompt:

defadd(a: float, b: float = 2.5) -> float:
                if a < 0:
                    raise ValueError("Provide a positive number for a")
                try:
                    return a + b
                except Exception as e:
                    raise (e)
          

L'output du modèle semble relativement correct sur un plan sémantique mais incomplet et mal formaté:

defadd(a: float, b: float = 2.5) -> float:
                """
                This function takes two floating point numbers as input and returns their sum.

                Args:
                a (float): The first number to add.
                b (float, optional): The second number to add. Defaults to 2.5.

                Returns:
                float: The sum of the two input numbers.

                Raises:
                ValueError: If the input number a is negative.
                """if a < 0:
                    raise ValueError("Provide a positive number for a")
                try:
                    return a + b
                except Exception as e:
                    raise (e)

                # Example usage:
                result = add(3.14, 1.27)
                print(result)  # Output: 4.41

Noter que le modèle retourne l'intégralité de la fonction et non juste le docstring comme demandé en instruction. Il n'a pas non plus "compris" que nous souhaitons avoir l'exemple dans le docstring, et non pas en tant que code supplémentaire. Pour corriger ces erreurs nous devons améliorer le template, en montrant au modèle un exemple de l'ouptut que nous souhaitons obtenir.

Few shots prompting

Pour affiner ce que nous souhaitons obtenir du modèle, nous allons lui fournir des exemples (des "shots" dans la terminologie standard). Reprenons la fonction ci-dessus pour créer un prompt one shot qui définira le type d'ouput précis que nous souhaitons: ici un docstring détaillé avec un exemple. Reprenons notre template pour y ajouter l'exemple:

"""<s>[INST] in Python create a detailed and helpful Google style docstring for this code:

              ```python
              def add(a: float, b: float = 2.5) -> float:
                  if a < 0:
                      raise ValueError("Provide a positive number for a")
                  try:
                      return a + b
                  except Exception as e:
                      raise (e)
              ```
              
              Important: always provide a short example in the docstring. Break lines at 86 characters maximum. Output only the docstring [/INST]
              
              ```python
              \"\"\"
              Sums two float numbers, but ensures the first number is non-negative. If the
              second number is not provided, it defaults to 2.5.
              
              Args:
                  a (float): The first number to be added. Must be a non-negative float.
                  b (float, optional): The second number to be added. Defaults to 2.5. Can be
                      any float.
              
              Returns:
                  float: The sum of a and b.
              
              Raises:
                  ValueError: If `a` is negative.
                  Exception: Any unexpected error that might occur during addition.
              
              Example:
                  >>> add(2.5)
                  5.0
                  >>> add(2.5, 3.5)
                  6.0
                  >>> add(-1.0)
                  ValueError: Provide a positive number for a
              \"\"\"
              ```</s>
              [INST] in Python create a detailed and helpful Google style docstring for this code:
              
              ```python
              {prompt}
              ```
              
              Important: always provide a short example in the docstring. Break lines at 86 characters maximum. Output only the docstring [/INST]"""

Retestons maintenant le template, avec ce code passé en prompt:

defdivide(x, y):
                try:
                    result = x / y
                except ZeroDivisionError:
                    print("division by zero!")
                else:
                    print("result is", result)
                finally:
                    print("executing finally clause")
          

Le résultat est maintenant tout à fait satisfaisant: le modèle retourne uniquement un bloc Python markdown avec le docstring, comme demandé:

"""
            Divides two numbers and handles division by zero.

            Args:
                x (float): The first number to be divided.
                y (float): The second number to be divided.

            Returns:
                None: No return value is produced.

            Raises:
                ZeroDivisionError: If `y` is zero.

            Example:
                >>> divide(10, 2)
                result is 5.0
                executing finally clause
                >>> divide(10, 0)
                division by zero!
                executing finally clause
            """

Le docstring est correctement formaté, et l'exemple inclus dedans comme demandé. Le shot présenté au modèle a considérablement amélioré la précision de l'output. Important: toujours bien relire et vérifier le caractère correct de ce qui est produit par le modèle: des imprécisions ou erreurs peuvent parfois apparaître, selon de le phénomène dit d'hallucination notamment, où le modèle invente des choses.

Retrouvez le code de cet exemple ainsi que d'autres dans la section exemples de la lib LocalLm.

Conclusion

Les progrès récents dans le domaine de l'IA open source permettent d'envisager une utilisation réaliste de modèles de taille relativement modestes pour certaines tâches, pouvant donc tourner sur des configurations locales moyennes. Noter que la présence d'un GPU permet d'accélerer notablement la vitesse du modèle, notamment l'ingestion des prompts, lente avec du CPU seul.

La manière d'interroger le modèle, notamment le format du template utilisé demeure essentielle pour obtenir un output de qualité optimale. Un format de template inadapté au modèle conduira à des résultats d'inférence moins bons, voire fera complètement dérailler le modèle.

Des compétences avancées en prompt engineering permettent de tirer le meilleur des modèles open source disponibles, et d'envisager une utilisation profitable. De nombreuses techniques de prompting existent, suivant la nature de tâches à accomplir: few shots prompting comme montré ici, mais aussi tree of toughts ou chain of toughts par exemple, certaines reposant sur des agents collaboratifs impliquant plusieurs modèles spécialisés.

Articles associés

LocalLm : une IA locale en Python

Face à ChatGPT et autres large language models (LLM) fermés, des alternatives open source se développent. De plus en plus de models open source font leur apparition : les Llama, Mistral et autres Falcon et consorts. Cet écosystème émergent permet d'envisager une utilisation interne en entreprise, hors plateformes géantes. Même si les coûts et la technologie restent incertains et en constante évolution, une telle approche permet de préserver la confidentialité des données, en ne les envoyant pas sur des plateformes externes.

Utiliser une IA locale dans Django

Ceci est une introduction l'IA prend de plus en plus de place, de plus en plus vite. Depuis quelques temps a émergé un mouvement open source autour notamment de Llama.cpp.

Django et IA locale : connecter Langchain