import os
from dotenv import load_dotenv, find_dotenv
from langchain import HuggingFaceHub
from langchain import PromptTemplate, LLMChain, OpenAI
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.chains.summarize import load_summarize_chain
from langchain.document_loaders import YoutubeLoader
import textwrap

# --------------------------------------------------------------
# Load the HuggingFaceHub API token from the .env file
# --------------------------------------------------------------

load_dotenv(find_dotenv())
HUGGINGFACEHUB_API_TOKEN = os.environ["HUGGINGFACEHUB_API_TOKEN"]


# --------------------------------------------------------------
# Load the LLM model from the HuggingFaceHub
# --------------------------------------------------------------

repo_id = "tiiuae/falcon-7b-instruct"  # See https://huggingface.co/models?pipeline_tag=text-generation&sort=downloads for some other options
falcon_llm = HuggingFaceHub(
    repo_id=repo_id, model_kwargs={"temperature": 0.1, "max_new_tokens": 500}
)


# --------------------------------------------------------------
# Create a PromptTemplate and LLMChain
# --------------------------------------------------------------
template = """Question: {question}

Answer: Let's think step by step."""

prompt = PromptTemplate(template=template, input_variables=["question"])
llm_chain = LLMChain(prompt=prompt, llm=falcon_llm)


# --------------------------------------------------------------
# Run the LLMChain
# --------------------------------------------------------------

question = "How do I make a sandwich?"
response = llm_chain.run(question)
wrapped_text = textwrap.fill(
    response, width=100, break_long_words=False, replace_whitespace=False
)
print(wrapped_text)


# --------------------------------------------------------------
# Load a video transcript from YouTube
# --------------------------------------------------------------

video_url = "https://www.youtube.com/watch?v=riXpu1tHzl0"
loader = YoutubeLoader.from_youtube_url(video_url)
transcript = loader.load()

text_splitter = RecursiveCharacterTextSplitter(chunk_size=3000)
docs = text_splitter.split_documents(transcript)

# --------------------------------------------------------------
# Summarization with LangChain
# --------------------------------------------------------------

# Add map_prompt and combine_prompt to the chain for custom summarization
chain = load_summarize_chain(falcon_llm, chain_type="map_reduce", verbose=True)
print(chain.llm_chain.prompt.template)
print(chain.combine_document_chain.llm_chain.prompt.template)

# --------------------------------------------------------------
# Test the Falcon model with text summarization
# --------------------------------------------------------------

output_summary = chain.run(docs)
wrapped_text = textwrap.fill(
    output_summary, width=100, break_long_words=False, replace_whitespace=False
)
print(wrapped_text)


# --------------------------------------------------------------
# Load an OpenAI model for comparison
# --------------------------------------------------------------

openai_llm = OpenAI(
    model_name="text-davinci-003", temperature=0.1, max_tokens=500
)  # max token length is 4097
chain = load_summarize_chain(openai_llm, chain_type="map_reduce", verbose=True)
output_summary = chain.run(docs)
wrapped_text = textwrap.fill(
    output_summary, width=100, break_long_words=False, replace_whitespace=False
)
print(wrapped_text)

langchain.chains.summarize.load_summarize_chain langchain.LLMChain langchain.text_splitter.RecursiveCharacterTextSplitter langchain.OpenAI langchain.document_loaders.YoutubeLoader.from_youtube_url langchain.HuggingFaceHub langchain.PromptTemplate
[((955, 1048), 'langchain.HuggingFaceHub', 'HuggingFaceHub', ([], {'repo_id': 'repo_id', 'model_kwargs': "{'temperature': 0.1, 'max_new_tokens': 500}"}), "(repo_id=repo_id, model_kwargs={'temperature': 0.1,\n    'max_new_tokens': 500})\n", (969, 1048), False, 'from langchain import HuggingFaceHub\n'), ((1305, 1368), 'langchain.PromptTemplate', 'PromptTemplate', ([], {'template': 'template', 'input_variables': "['question']"}), "(template=template, input_variables=['question'])\n", (1319, 1368), False, 'from langchain import PromptTemplate, LLMChain, OpenAI\n'), ((1381, 1420), 'langchain.LLMChain', 'LLMChain', ([], {'prompt': 'prompt', 'llm': 'falcon_llm'}), '(prompt=prompt, llm=falcon_llm)\n', (1389, 1420), False, 'from langchain import PromptTemplate, LLMChain, OpenAI\n'), ((1662, 1750), 'textwrap.fill', 'textwrap.fill', (['response'], {'width': '(100)', 'break_long_words': '(False)', 'replace_whitespace': '(False)'}), '(response, width=100, break_long_words=False,\n    replace_whitespace=False)\n', (1675, 1750), False, 'import textwrap\n'), ((2012, 2053), 'langchain.document_loaders.YoutubeLoader.from_youtube_url', 'YoutubeLoader.from_youtube_url', (['video_url'], {}), '(video_url)\n', (2042, 2053), False, 'from langchain.document_loaders import YoutubeLoader\n'), ((2098, 2145), 'langchain.text_splitter.RecursiveCharacterTextSplitter', 'RecursiveCharacterTextSplitter', ([], {'chunk_size': '(3000)'}), '(chunk_size=3000)\n', (2128, 2145), False, 'from langchain.text_splitter import RecursiveCharacterTextSplitter\n'), ((2440, 2511), 'langchain.chains.summarize.load_summarize_chain', 'load_summarize_chain', (['falcon_llm'], {'chain_type': '"""map_reduce"""', 'verbose': '(True)'}), "(falcon_llm, chain_type='map_reduce', verbose=True)\n", (2460, 2511), False, 'from langchain.chains.summarize import load_summarize_chain\n'), ((2841, 2935), 'textwrap.fill', 'textwrap.fill', (['output_summary'], {'width': '(100)', 'break_long_words': '(False)', 'replace_whitespace': '(False)'}), '(output_summary, width=100, break_long_words=False,\n    replace_whitespace=False)\n', (2854, 2935), False, 'import textwrap\n'), ((3142, 3212), 'langchain.OpenAI', 'OpenAI', ([], {'model_name': '"""text-davinci-003"""', 'temperature': '(0.1)', 'max_tokens': '(500)'}), "(model_name='text-davinci-003', temperature=0.1, max_tokens=500)\n", (3148, 3212), False, 'from langchain import PromptTemplate, LLMChain, OpenAI\n'), ((3255, 3326), 'langchain.chains.summarize.load_summarize_chain', 'load_summarize_chain', (['openai_llm'], {'chain_type': '"""map_reduce"""', 'verbose': '(True)'}), "(openai_llm, chain_type='map_reduce', verbose=True)\n", (3275, 3326), False, 'from langchain.chains.summarize import load_summarize_chain\n'), ((3375, 3469), 'textwrap.fill', 'textwrap.fill', (['output_summary'], {'width': '(100)', 'break_long_words': '(False)', 'replace_whitespace': '(False)'}), '(output_summary, width=100, break_long_words=False,\n    replace_whitespace=False)\n', (3388, 3469), False, 'import textwrap\n'), ((541, 554), 'dotenv.find_dotenv', 'find_dotenv', ([], {}), '()\n', (552, 554), False, 'from dotenv import load_dotenv, find_dotenv\n')]