Update chain.py

chain.py

@@ -1,6 +1,7 @@
 import json
 import os
 import pathlib
+import pickle
 from typing import Dict, List, Tuple
 
 import weaviate
@@ -17,12 +18,9 @@ from langchain.prompts.example_selector import \
 from langchain.vectorstores import FAISS, Weaviate
 from pydantic import BaseModel
 
-
-
-
 class CustomChain(Chain, BaseModel):
 
-    vstore: Weaviate
+    vstore: FAISS
     chain: BaseCombineDocumentsChain
     key_word_extractor: Chain
 
@@ -52,13 +50,7 @@ class CustomChain(Chain, BaseModel):
         return {"answer": answer}
 
 
-def get_new_chain1(vectorstore) -> Chain:
-    WEAVIATE_URL = os.environ["WEAVIATE_URL"]
-    client = weaviate.Client(
-        url=WEAVIATE_URL,
-        additional_headers={"X-OpenAI-Api-Key": os.environ["OPENAI_API_KEY"]},
-    )
-
+def get_new_chain1(vectorstore, rephraser_llm, final_output_llm) -> Chain:
     _eg_template = """## Example:
 
     Chat History:
@@ -70,40 +62,39 @@ def get_new_chain1(vectorstore) -> Chain:
         input_variables=["chat_history", "question", "answer"],
     )
 
-    _prefix = """Given the following conversation and a follow up question, rephrase the follow up question to be a standalone question. You should assume that the question is related to LangChain."""
+    _prefix = """Given the following conversation and a follow up question, rephrase the follow up question to be a standalone question. You should assume that the question is related to Hugging Face Code."""
     _suffix = """## Example:
 
     Chat History:
     {chat_history}
     Follow Up Input: {question}
     Standalone question:"""
-    eg_store = Weaviate(
-        client,
-        "Rephrase",
-        "content",
-        attributes=["question", "answer", "chat_history"],
-    )
-    example_selector = SemanticSimilarityExampleSelector(vectorstore=eg_store, k=4)
+
+    #### LOAD VSTORE WITH REPHRASE EXAMPLES
+    with open("rephrase_eg.pkl", 'rb') as f:
+        rephrase_example_selector = pickle.load(f)
+    
     prompt = FewShotPromptTemplate(
         prefix=_prefix,
         suffix=_suffix,
-        example_selector=example_selector,
+        example_selector=rephrase_example_selector,
         example_prompt=_eg_prompt,
         input_variables=["question", "chat_history"],
     )
-    llm = OpenAI(temperature=0, model_name="text-davinci-003")
-    key_word_extractor = LLMChain(llm=llm, prompt=prompt)
+
+    key_word_extractor = LLMChain(llm=rephraser_llm, prompt=prompt)
 
     EXAMPLE_PROMPT = PromptTemplate(
         template=">Example:\nContent:\n---------\n{page_content}\n----------\nSource: {source}",
         input_variables=["page_content", "source"],
     )
-    template = """You are an AI assistant for the open source library LangChain. The documentation is located at https://langchain.readthedocs.io.
-You are given the following extracted parts of a long document and a question. Provide a conversational answer with a hyperlink to the documentation.
+
+    template = """You are an AI assistant for the open source transformers library provided by Hugging Face. The documentation is located at https://huggingface.co/docs/transformers.
+You are given the following extracted parts of a long document and a question. Provide a conversational answer with a hyperlink to the documentation. Do NOT add .html to the end of links.
 You should only use hyperlinks that are explicitly listed as a source in the context. Do NOT make up a hyperlink that is not listed.
 If the question includes a request for code, provide a code block directly from the documentation.
 If you don't know the answer, just say "Hmm, I'm not sure." Don't try to make up an answer.
-If the question is not about LangChain, politely inform them that you are tuned to only answer questions about LangChain.
+If the question is not about Hugging Face Transformers, politely inform them that you are tuned to only answer questions about Transformers.
 Question: {question}
 =========
 {context}
@@ -111,10 +102,11 @@ Question: {question}
 Answer in Markdown:"""
     PROMPT = PromptTemplate(template=template, input_variables=["question", "context"])
     doc_chain = load_qa_chain(
-        OpenAI(temperature=0, model_name="text-davinci-003", max_tokens=-1),
+        final_output_llm,
         chain_type="stuff",
         prompt=PROMPT,
         document_prompt=EXAMPLE_PROMPT,
+        verbose=True
     )
     return CustomChain(chain=doc_chain, vstore=vectorstore, key_word_extractor=key_word_extractor)