feat: add history

main.py

@@ -12,6 +12,7 @@ import subprocess
 from langchain.chains import RetrievalQA
 from langchain.embeddings import HuggingFaceInstructEmbeddings
 from langchain.prompts import PromptTemplate
+from langchain.memory import ConversationBufferMemory
 
 # from langchain.embeddings import HuggingFaceEmbeddings
 from run_localGPT import load_model
@@ -55,11 +56,13 @@ Always answer in the most helpful and safe way possible.
 If you don't know the answer to a question, just say that you don't know, don't try to make up an answer, don't share false information.
 Use 15 sentences maximum. Keep the answer as concise as possible.
 Always say "thanks for asking!" at the end of the answer.
-{context}
+Context: {history} \n {context}
 Question: {question}
-Helpful Answer:"""
+"""
 
-QA_CHAIN_PROMPT = PromptTemplate.from_template(template)
+memory = ConversationBufferMemory(input_key="question", memory_key="history")
+
+QA_CHAIN_PROMPT = PromptTemplate.from_template(input_variables=["history", "context", "question"], template=template)
 
 QA = RetrievalQA.from_chain_type(
     llm=LLM,
@@ -68,6 +71,7 @@ QA = RetrievalQA.from_chain_type(
     return_source_documents=SHOW_SOURCES,
     chain_type_kwargs={
         "prompt": QA_CHAIN_PROMPT,
+        "memory": memory
     },
 )
 
@@ -118,7 +122,6 @@ def run_ingest_route():
         )
 
         RETRIEVER = DB.as_retriever()
-        prompt, memory = get_prompt_template(promptTemplate_type="llama", history=True)
 
         QA = RetrievalQA.from_chain_type(
             llm=LLM,
@@ -127,12 +130,11 @@ def run_ingest_route():
             return_source_documents=SHOW_SOURCES,
             chain_type_kwargs={
                 "prompt": QA_CHAIN_PROMPT,
+                "memory": memory
             },
         )
 
-        response = "Script executed successfully: {}".format(result.stdout.decode("utf-8"))
-
-        return {"response": response}
+        return {"response": "The training was successfully completed"}
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error occurred: {str(e)}")
 

run_localGPT.py

@@ -79,7 +79,7 @@ def load_model(device_type, model_id, model_basename=None, LOGGING=logging):
 
     # Create a pipeline for text generation
 
-    streamer = TextStreamer(tokenizer, skip_prompt=True)
+    streamer = TextStreamer(tokenizer)
 
     pipe = pipeline(
         "text-generation",