Some minor adjustments.

souradipp76 · Apr 29, 2024 · 5bf7023 · 5bf7023
1 parent 3da584c
commit 5bf7023
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 6 deletions.
diff --git a/doc_generator/autodoc/docs/data/docstore.json b/doc_generator/autodoc/docs/data/docstore.json
diff --git a/doc_generator/autodoc/docs/data/hnswlib.index b/doc_generator/autodoc/docs/data/hnswlib.index
diff --git a/doc_generator/query/__init__.py b/doc_generator/query/__init__.py
@@ -40,9 +40,9 @@ def query(repo_config: AutodocRepoConfig, user_confg: AutodocUserConfig):
         print('Thinking...')
         try:
             response = chain.invoke({'question': question, 'chat_history': chat_history})
-            chat_history.append((question, response['text']))
+            chat_history.append((question, response['answer']))
             print('\n\nMarkdown:\n')
-            print(markdown(response['text']))
+            print(markdown(response['answer']))
         except Exception as error:
             print(f"Something went wrong: {error}")
             traceback.print_exc()
diff --git a/doc_generator/utils/LLMUtils.py b/doc_generator/utils/LLMUtils.py
@@ -9,8 +9,8 @@
 
 def get_chat_model(model_name: str, model_kwargs):
     config = AutoConfig.from_pretrained(model_name)
-    config.quantization_config["disable_exllama"] = True
-    config.quantization_config["exllama_config"] = {"version":2}
+    # config.quantization_config["disable_exllama"] = True
+    # config.quantization_config["exllama_config"] = {"version":2}
     tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
     model = AutoModelForCausalLM.from_pretrained(
                 model_name,
@@ -117,7 +117,7 @@ def total_index_cost_estimate(model):
 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: