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.

“Jonathan Sugarman died at Camp 2 after he began to feel unwell,” Pasang Sherpa told CNN on Tuesday.

Seattle-based Sugarman was part of an expedition arranged by Washington state-based International Mountain Guides (IMG) with Beyul Adventure handling the local logistics.

Sherpa added that “his body remains at Camp 2 with the rest of the climbing team.”

This comes after Nepal has issued permits for a record 463 climbers by April 26, for this spring season’s expeditions to Mount Everest.

Following Sugarman’s death, the Embassy of the United States issued a statement. “We can confirm Dr. Jonathan Sugarman passed away while climbing Mt. Everest Monday May 1,” it said. “Our deepest sympathies go out to his family and friends.

“The Embassy is in contact with Dr. Sugarman’s family and with local authorities. Out of respect for the family’s privacy, we cannot comment further,” read a statement sent to CNN by an Embassy spokesperson.
'''
    
    return s

def mainSimple():
    print("\nIn 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", return_full_text=True)
    inputText = getText()
    resp = generate_text_pipline(inputText)
    respStr = resp[0]["generated_text"] 
    
    print(f'\nInput: {inputText}')
    print(f'\nResponse: {respStr}')
    
    print("\nAll Done!\n")

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("\nIn 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'\nInput-Context: {context}')
    print(f'\nInput-Instr: {instr}')
    print(f'\nResponse: {resp}')
    
    print("\nAll Done!\n") 

if __name__ == "__main__":
    mainContext()