Some fixes in llama2

doc_generator/main.py

@@ -34,7 +34,7 @@ def main():  # pragma: no cover
         "root": "./autodoc",
         "repository_url": "https://github.com/context-labs/autodoc",
         "output": "doc_generator/autodoc",
-        "llms": ["gpt-3.5-turbo"],
+        "llms": ["TheBloke/Llama-2-7B-Chat-GPTQ"],
         "ignore": [
             ".*",
             "*package-lock.json",

doc_generator/types.py

@@ -1,6 +1,7 @@
 from enum import Enum
 from typing import List, Callable, Optional
 from langchain_openai import ChatOpenAI
+from langchain_experimental.chat_models import Llama2Chat
 
 class LLMModels(str, Enum):
     GPT3 = "gpt-3.5-turbo"
@@ -101,7 +102,7 @@ def __init__(self, input_path: str, project_name: str, process_file: Optional[Ca
 
 class LLMModelDetails:
     def __init__(self, name: LLMModels, input_cost_per_1k_tokens: float, output_cost_per_1k_tokens: float,
-                 max_length: int, llm: LLama2Chat, input_tokens: int, output_tokens: int, succeeded: int,
+                 max_length: int, llm: Llama2Chat, input_tokens: int, output_tokens: int, succeeded: int,
                  failed: int, total: int):
         self.name = name
         self.input_cost_per_1k_tokens = input_cost_per_1k_tokens

doc_generator/utils/LLMUtils.py

@@ -1,10 +1,10 @@
 import os
 import torch
 
-from langchain.embeddings import HuggingFaceEmbeddings
+from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_openai import OpenAIEmbeddings, ChatOpenAI
 from langchain_experimental.chat_models import Llama2Chat
-from langchain import HuggingFacePipeline
+from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
 from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig, pipeline, BitsAndBytesConfig
 from doc_generator.types import LLMModelDetails, LLMModels
 
@@ -53,18 +53,12 @@
         llm=Llama2Chat(llm=HuggingFacePipeline(pipeline=pipeline(
             "text-generation",
             model=AutoModelForCausalLM.from_pretrained(
-                LLMModels.LLAMA2_7B_CHAT_GPTQ,
-                quantization_config=BitsAndBytesConfig(
-                    load_in_4bit=True,
-                    bnb_4bit_quant_type='nf4',
-                    bnb_4bit_use_double_quant=True,
-                    bnb_4bit_compute_dtype=torch.bfloat16
-                ),
+                LLMModels.LLAMA2_7B_CHAT_GPTQ.value,
                 torch_dtype=torch.float16,
                 trust_remote_code=True,
                 device_map="auto"
             ),
-            tokenizer=AutoTokenizer.from_pretrained(LLMModels.LLAMA2_7B_CHAT_GPTQ, use_fast=True),
+            tokenizer=AutoTokenizer.from_pretrained(LLMModels.LLAMA2_7B_CHAT_GPTQ.value, use_fast=True),
             generation_config=AutoConfig.from_pretrained(
                 LLMModels.LLAMA2_7B_CHAT_GPTQ,
             ),
@@ -77,6 +71,7 @@
     )
 }
 
+
 def print_model_details(models):
     output = []
     for model_details in models.values():
@@ -104,6 +99,7 @@ def print_model_details(models):
     for item in all_results:
         print(item)
 
+
 def total_index_cost_estimate(model):
     total_cost = sum(
         (model.input_tokens / 1000) * model.input_cost_per_1k_tokens +
@@ -112,26 +108,22 @@ def total_index_cost_estimate(model):
     )
     return total_cost
 
-def get_embeddings(model):
-    if model == LLMModels.LLAMA2_7B_CHAT_GPTQ:
+
+def get_embeddings(model:str):
+    if model == LLMModels.LLAMA2_7B_CHAT_GPTQ.value:
         return HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2",
-                                     model_kwargs={"device": "cuda"},
+                                     # model_kwargs={"device": "cuda"},
                                      encode_kwargs={"normalize_embeddings": True},
                                      )
     else:
         return OpenAIEmbeddings()
 
-def get_chat_model(model_name, model_kwargs):
+
+def get_chat_model(model_name: str, model_kwargs):
     return Llama2Chat(llm=HuggingFacePipeline(pipeline=pipeline(
             "text-generation",
             model=AutoModelForCausalLM.from_pretrained(
                 model_name,
-                quantization_config=BitsAndBytesConfig(
-                    load_in_4bit=True,
-                    bnb_4bit_quant_type='nf4',
-                    bnb_4bit_use_double_quant=True,
-                    bnb_4bit_compute_dtype=torch.bfloat16
-                ),
                 torch_dtype=torch.float16,
                 trust_remote_code=True,
                 device_map="auto"

doc_generator/utils/createChatChain.py

@@ -11,30 +11,58 @@
                 "Chat History:\n{chat_history}\nFollow Up Input: {question}\nStandalone question:"
 )
 
-def make_qa_prompt(project_name, repository_url, content_type, chat_prompt, target_audience):
-    additional_instructions = f"\nHere are some additional instructions for answering questions about {content_type}:\n{chat_prompt}" if chat_prompt else ""
-    return PromptTemplate.from_template(
-        f"You are an AI assistant for a software project called {project_name}. You are trained on all the {content_type} that makes up this project.\n"
-        f"The {content_type} for the project is located at {repository_url}.\n"
-        "You are given a repository which might contain several modules and each module will contain a set of files.\n"
-        "Look at the source code in the repository and you have to geneate a readme.md file following some template given below. If you use any hyperlinks, they should link back to the github repository shared with you.\n"
-        "You should only use hyperlinks that are explicitly listed in the context. Do NOT make up a hyperlink that is not listed.\n"
+# def make_qa_prompt(project_name, repository_url, content_type, chat_prompt, target_audience):
+#     additional_instructions = f"\nHere are some additional instructions for answering questions about {content_type}:\n{chat_prompt}" if chat_prompt else ""
+#     return PromptTemplate.from_template(
+#         f"You are an AI assistant for a software project called {project_name}. You are trained on all the {content_type} that makes up this project.\n"
+#         f"The {content_type} for the project is located at {repository_url}.\n"
+#         "You are given a repository which might contain several modules and each module will contain a set of files.\n"
+#         "Look at the source code in the repository and you have to geneate a readme.md file following some template given below. If you use any hyperlinks, they should link back to the github repository shared with you.\n"
+#         "You should only use hyperlinks that are explicitly listed in the context. Do NOT make up a hyperlink that is not listed.\n"
 
-        "Assume the reader is a {target_audience} but is not deeply familiar with {project_name}.\n"
-        "Assume the reader does not know anything about how the project is structured or which folders/files do what and what functions are written in which files and what these functions do.\n"
-        "If you don't know how to fill up the readme.md file in one of its sections, leave that part blank. Don't try to make up an answer.\n"
-        "Do not include information that is not directly relevant to repository, even though the names of the functions might be common or is frequently used in several other places.\n"
-        "Now lets start describing how the readme.md file will be structured.\n"
-        "The first section will be Installation. Here provide a list of packages from the requirements.txt folder in the repository that needs to be installed. Mention what versions of those packages need to be installed. Also add the commands that need to be put in the terminal to install those packages. For instance, for installing a py-package, provide a command pip install py-package. If there is no requirements.txt or similar folder in the repository, then find out what frameworks and packages have been imported in all the files after going through the code provide their names and the required versions that need to be installed. Remind the user that it is usually best-practice in Python projects to install into a sandboxed virtual environment, This will be locked to a specific Python version and contain only the Python libraries that you install into it, so that your Python projects do not get affected.\n"
-        "The second section will be Usage. Here provide a list of commands that need to be run in the terminal to use various features of the project. For instance, if the project has a command called run, then provide a command to run that command. Go through various files and various modules and after reading each function, provide an example usage of that function. Write two lines about each function, what the functionality of that function is, what paramters it take as input, what is outputs, if it is dependent on the output of any other function in the whole repository, what other functions call/use this function, and finally provide a toy example of the usage of the function. Do this for every function in all files that exist in the repository. Structure them in the same way the repository has been structured. Finally provide some run commands on how to run the main function in the terminal.\n"
-        "The third section will be Development. Here provide a list of commands that need to be run in the terminal to develop. Comment here on how to make format, make lint, check types and generate tests for the code that has been written in the files of the repository. Try to create unit tests and write tests for functions that are called/ used by many other functions, to test that they work correctly. Write some commands on how to run those tests.\n"
-        "The fourth section will be Conclusion. Here provide a general idea of what the project does and how it works. Here you can also provide how the user can geenrate a github.io page for the repository documentation using the four sections that you generated above in the readme.md file. Tell how to create a workflow and yml file and github pages can be used to create a documentation for the project. Put the readme.md file in the repository.\n"
+#         f"Assume the reader is a {target_audience} but is not deeply familiar with {project_name}.\n"
+#         "Assume the reader does not know anything about how the project is structured or which folders/files do what and what functions are written in which files and what these functions do.\n"
+#         "If you don't know how to fill up the readme.md file in one of its sections, leave that part blank. Don't try to make up an answer.\n"
+#         "Do not include information that is not directly relevant to repository, even though the names of the functions might be common or is frequently used in several other places.\n"
+#         "Now lets start describing how the readme.md file will be structured.\n"
+#         "The first section will be Installation. Here provide a list of packages from the requirements.txt folder in the repository that needs to be installed. Mention what versions of those packages need to be installed. Also add the commands that need to be put in the terminal to install those packages. For instance, for installing a py-package, provide a command pip install py-package. If there is no requirements.txt or similar folder in the repository, then find out what frameworks and packages have been imported in all the files after going through the code provide their names and the required versions that need to be installed. Remind the user that it is usually best-practice in Python projects to install into a sandboxed virtual environment, This will be locked to a specific Python version and contain only the Python libraries that you install into it, so that your Python projects do not get affected.\n"
+#         "The second section will be Usage. Here provide a list of commands that need to be run in the terminal to use various features of the project. For instance, if the project has a command called run, then provide a command to run that command. Go through various files and various modules and after reading each function, provide an example usage of that function. Write two lines about each function, what the functionality of that function is, what paramters it take as input, what is outputs, if it is dependent on the output of any other function in the whole repository, what other functions call/use this function, and finally provide a toy example of the usage of the function. Do this for every function in all files that exist in the repository. Structure them in the same way the repository has been structured. Finally provide some run commands on how to run the main function in the terminal.\n"
+#         "The third section will be Development. Here provide a list of commands that need to be run in the terminal to develop. Comment here on how to make format, make lint, check types and generate tests for the code that has been written in the files of the repository. Try to create unit tests and write tests for functions that are called/ used by many other functions, to test that they work correctly. Write some commands on how to run those tests.\n"
+#         "The fourth section will be Conclusion. Here provide a general idea of what the project does and how it works. Here you can also provide how the user can geenrate a github.io page for the repository documentation using the four sections that you generated above in the readme.md file. Tell how to create a workflow and yml file and github pages can be used to create a documentation for the project. Put the readme.md file in the repository.\n"
 
-        f"{additional_instructions}\n"
-        "Question: {question}\n\n"
-        "Context:\n{context}\n\n"
-        "Answer in Markdown:"
-    )
+#         f"{additional_instructions}\n"
+#         "Question: {question}\n\n"
+#         "Context:\n{context}\n\n"
+#         "Answer in Markdown:"
+#     )
+
+
+def make_qa_prompt(project_name, repository_url, content_type, chat_prompt, target_audience):
+    additional_instructions = f"\nHere are some additional instructions for answering questions about {content_type}:\n{chat_prompt}" if chat_prompt else ""
+    template = f"""<s>[INST]You are an AI assistant for a software project called {project_name}. You are trained on all the {content_type} that makes up this project.\n"
+    f"The {content_type} for the project is located at {repository_url}.\n"
+    "You are given a repository which might contain several modules and each module will contain a set of files.\n"
+    "Look at the source code in the repository and you have to geneate a readme.md file following some template given below. If you use any hyperlinks, they should link back to the github repository shared with you.\n"
+    "You should only use hyperlinks that are explicitly listed in the context. Do NOT make up a hyperlink that is not listed.\n"
+
+    f"Assume the reader is a {target_audience} but is not deeply familiar with {project_name}.\n"
+    "Assume the reader does not know anything about how the project is structured or which folders/files do what and what functions are written in which files and what these functions do.\n"
+    "If you don't know how to fill up the readme.md file in one of its sections, leave that part blank. Don't try to make up an answer.\n"
+    "Do not include information that is not directly relevant to repository, even though the names of the functions might be common or is frequently used in several other places.\n"
+    "Now lets start describing how the readme.md file will be structured.\n"
+    "The first section will be Installation. Here provide a list of packages from the requirements.txt folder in the repository that needs to be installed. Mention what versions of those packages need to be installed. Also add the commands that need to be put in the terminal to install those packages. For instance, for installing a py-package, provide a command pip install py-package. If there is no requirements.txt or similar folder in the repository, then find out what frameworks and packages have been imported in all the files after going through the code provide their names and the required versions that need to be installed. Remind the user that it is usually best-practice in Python projects to install into a sandboxed virtual environment, This will be locked to a specific Python version and contain only the Python libraries that you install into it, so that your Python projects do not get affected.\n"
+    "The second section will be Usage. Here provide a list of commands that need to be run in the terminal to use various features of the project. For instance, if the project has a command called run, then provide a command to run that command. Go through various files and various modules and after reading each function, provide an example usage of that function. Write two lines about each function, what the functionality of that function is, what paramters it take as input, what is outputs, if it is dependent on the output of any other function in the whole repository, what other functions call/use this function, and finally provide a toy example of the usage of the function. Do this for every function in all files that exist in the repository. Structure them in the same way the repository has been structured. Finally provide some run commands on how to run the main function in the terminal.\n"
+    "The third section will be Development. Here provide a list of commands that need to be run in the terminal to develop. Comment here on how to make format, make lint, check types and generate tests for the code that has been written in the files of the repository. Try to create unit tests and write tests for functions that are called/ used by many other functions, to test that they work correctly. Write some commands on how to run those tests.\n"
+    "The fourth section will be Conclusion. Here provide a general idea of what the project does and how it works. Here you can also provide how the user can geenrate a github.io page for the repository documentation using the four sections that you generated above in the readme.md file. Tell how to create a workflow and yml file and github pages can be used to create a documentation for the project. Put the readme.md file in the repository.\n"
+
+    f"{additional_instructions}\n"
+    "Question: {question}\n\n"
+    "Context:\n{context}\n\n"
+    "Answer in Markdown:"
+    "[/INST]"""
+
+    prompt = PromptTemplate(template=template, input_variables=["question", "context"])
+    return prompt
 
 def make_chain(project_name, repository_url, content_type, chat_prompt, target_audience, vectorstore, llms, on_token_stream=None):
     llm = llms[1] if len(llms) > 1 else llms[0]
@@ -53,9 +81,7 @@ def make_chain(project_name, repository_url, content_type, chat_prompt, target_a
         #                 "frequency_penalty": 0.0,
         #                 "presence_penalty": 0.0,
         #                }),
-        llm=get_chat_model(llm, {
-            "temperature": 0.2
-        }),
+        llm=get_chat_model(llm, {"temperature": 0.2}),
         prompt=qa_prompt
     )