modify tester to handle context-based instruction mode.

tester.py

@@ -3,6 +3,9 @@ import pathlib
 import torch
 from transformers import pipeline
 
+from langchain import PromptTemplate, LLMChain
+from langchain.llms import HuggingFacePipeline
+
 def getText():
     s = '''
     A US climber has died on his way to scale Mount Everest on Monday, according to an expedition organizer.
@@ -22,13 +25,13 @@ Following Sugarman’s death, the Embassy of the United States issued a statemen
     
     return s
 
-def main():
-    print("In main...")
+def mainSimple():
+    print("In main simple...")
     # print(pathlib.Path().resolve())
     # generate_text_pipline = pipeline(model="databricks/dolly-v2-3b", torch_dtype=torch.bfloat16, trust_remote_code=True, device_map="auto")
     # generate_text_pipline = pipeline(model="verseAI/databricks-dolly-v2-3b", torch_dtype=torch.bfloat16, trust_remote_code=True, device_map="auto")
     workingDir = pathlib.Path().resolve()
-    generate_text_pipline = pipeline(model=workingDir, torch_dtype=torch.bfloat16, trust_remote_code=True, device_map="auto")
+    generate_text_pipline = pipeline(model=workingDir, torch_dtype=torch.bfloat16, trust_remote_code=True, device_map="auto", return_full_text=True)
     inputText = getText()
     resp = generate_text_pipline(inputText)
     respStr = resp[0]["generated_text"] 
@@ -38,5 +41,47 @@ def main():
     
     print("All Done!")
 
+def getInstrAndContext():
+    context = '''George Washington (February 22, 1732[b] - December 14, 1799) was an American military officer, statesman,
+and Founding Father who served as the first president of the United States from 1789 to 1797.'''
+    # instr = '''When was George Washington president?'''
+    instr = '''What do you think?'''
+
+    return instr, context
+
+def mainContext():
+    print("In main context...")
+    workingDir = pathlib.Path().resolve()
+    generate_text_pipline = pipeline(model=workingDir, torch_dtype=torch.bfloat16, trust_remote_code=True, device_map="auto", return_full_text=True)
+
+    # template for an instrution with no input
+    prompt = PromptTemplate(
+        input_variables=["instruction"],
+        template="{instruction}")
+
+    # template for an instruction with input
+    prompt_with_context = PromptTemplate(
+        input_variables=["instruction", "context"],
+        template="{instruction}\n\nInput:\n{context}")
+
+    hf_pipeline = HuggingFacePipeline(pipeline=generate_text_pipline)
+
+    llm_chain = LLMChain(llm=hf_pipeline, prompt=prompt)
+    llm_context_chain = LLMChain(llm=hf_pipeline, prompt=prompt_with_context)
+
+    instr, context = getInstrAndContext()
+    resp = ''
+    
+    if(context and not context.isspace()):
+        resp = llm_context_chain.predict(instruction=instr, context=context).lstrip()
+    else:
+        resp = llm_chain.predict(instruction=instr).lstrip()
+
+    print(f'Input-Instr: {instr}')
+    print(f'Input-Context: {context}')
+    print(f'Response: {resp}')
+    
+    print("All Done!") 
+
 if __name__ == "__main__":
-    main()
+    mainContext()