###
# Elo based comparison of models
# https://chat.lmsys.org/?leaderboard
###

## 
# visual libraries gradio , could be streamlit as well or cl
##
import gradio as gr

##
# Libraries
# Langchain - https://python.langchain.com/docs/get_started/introduction.html
# Used for simplifiing calls, task
##
import langchain
import transformers


# https://huggingface.co/spaces/joyson072/LLm-Langchain/blob/main/app.py
from langchain.llms import HuggingFaceHub


# for the chain and prompt
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain, SimpleSequentialChain 

###################

llm = HuggingFaceHub(
    
    repo_id="google/flan-ul2", 
#    repo_id="google/flan-t5-small", 
    model_kwargs={"temperature":0.1,
                  "max_new_tokens":250})


# Chain 1: Generating a rephrased version of the user's question
template = """{question}\n\n"""
prompt_template = PromptTemplate(input_variables=["question"], template=template)
question_chain = LLMChain(llm=llm, prompt=prompt_template)

# Chain 2: Generating assumptions made in the statement
template = """Here is a statement:
    {statement}
    Make a bullet point list of the assumptions you made when producing the above statement.\n\n"""
prompt_template = PromptTemplate(input_variables=["statement"], template=template)
assumptions_chain = LLMChain(llm=llm, prompt=prompt_template)
assumptions_chain_seq = SimpleSequentialChain(
    chains=[question_chain, assumptions_chain], verbose=True
)

# Chain 3: Fact checking the assumptions
template = """Here is a bullet point list of assertions:
{assertions}
For each assertion, determine whether it is true or false. If it is false, explain why.\n\n"""
prompt_template = PromptTemplate(input_variables=["assertions"], template=template)
fact_checker_chain = LLMChain(llm=llm, prompt=prompt_template)
fact_checker_chain_seq = SimpleSequentialChain(
    chains=[question_chain, assumptions_chain, fact_checker_chain], verbose=True
)

# Final Chain: Generating the final answer to the user's question based on the facts and assumptions
template = """In light of the above facts, how would you answer the question '{}'""".format(
    "What is the capitol of the usa?"
#    user_question
)
template = """{facts}\n""" + template
prompt_template = PromptTemplate(input_variables=["facts"], template=template)
answer_chain = LLMChain(llm=llm, prompt=prompt_template)
overall_chain = SimpleSequentialChain(
    chains=[question_chain, assumptions_chain, fact_checker_chain, answer_chain],
    verbose=True,
)

#print(overall_chain.run("What is the capitol of the usa?"))

##################



#import model class and tokenizer
from transformers import BlenderbotTokenizer, BlenderbotForConditionalGeneration


###
# Definition of different purspose prompts
# https://huggingface.co/spaces/Chris4K/rlhf-arena/edit/main/app.py
####
def prompt_human_instruct(system_msg, history):
    return system_msg.strip() + "\n" + \
        "\n".join(["\n".join(["###Human: "+item[0], "###Assistant: "+item[1]])
        for item in history])


def prompt_instruct(system_msg, history):
    return system_msg.strip() + "\n" + \
        "\n".join(["\n".join(["### Instruction: "+item[0], "### Response: "+item[1]])
        for item in history])


def prompt_chat(system_msg, history):
    return system_msg.strip() + "\n" + \
        "\n".join(["\n".join(["USER: "+item[0], "ASSISTANT: "+item[1]])
        for item in history])


def prompt_roleplay(system_msg, history):
    return "<|system|>" + system_msg.strip() + "\n" + \
        "\n".join(["\n".join(["<|user|>"+item[0], "<|model|>"+item[1]])
        for item in history])


####
## Sentinent models
# https://huggingface.co/spaces/CK42/sentiment-model-comparison
# 1, 4 seem best for german
####
model_id_1 = "nlptown/bert-base-multilingual-uncased-sentiment"
model_id_2 = "microsoft/deberta-xlarge-mnli"
model_id_3 = "distilbert-base-uncased-finetuned-sst-2-english"
model_id_4 = "lordtt13/emo-mobilebert"
model_id_5 = "juliensimon/reviews-sentiment-analysis"
model_id_6 = "sbcBI/sentiment_analysis_model"
model_id_7 = "oliverguhr/german-sentiment-bert"

# https://colab.research.google.com/drive/1hrS6_g14EcOD4ezwSGlGX2zxJegX5uNX#scrollTo=NUwUR9U7qkld
#llm_hf_sentiment = HuggingFaceHub(
#    repo_id= model_id_7,
#    model_kwargs={"temperature":0.9 }
#)

from transformers import pipeline

# 
## Possible pipeline
#"['audio-classification', 'automatic-speech-recognition', 'conversational', 'depth-estimation', 'document-question-answering', 
#'feature-extraction', 'fill-mask', 'image-classification', 'image-segmentation', 'image-to-text', 'mask-generation', 'ner', 
#'object-detection', 'question-answering', 'sentiment-analysis', 'summarization', 'table-question-answering', 'text-classification', 
#'text-generation', 'text2text-generation', 'token-classification', 'translation', 'video-classification', 'visual-question-answering', 
#'vqa', 'zero-shot-audio-classification', 'zero-shot-classification', 'zero-shot-image-classification', 'zero-shot-object-detection', 
#'translation_XX_to_YY']"
##

sentiment_pipe = pipeline("sentiment-analysis", model=model_id_7)
#pipe = pipeline("translation", model="Helsinki-NLP/opus-mt-en-es")

def pipeline_predict_sentiment(text):
  sentiment_result = sentiment_pipe(text)
  print(sentiment_result)
  return sentiment_result


chat_pipe = pipeline("conversational")

def pipeline_predict_chat(text):
  sentiment_result = chat_pipe(text)
  print(sentiment_result)
  return sentiment_result


#['huggingface', 'models', 'spaces']
#sentiment = gr.load(model_id_7, src="huggingface")

#def sentiment (message):
#  sentiment_label = sentiment.predict(message)
#  print ( sentiment_label)
#  return sentiment_label

#sentiment_prompt = PromptTemplate(
#    input_variables=["text_input"],
#    template="Extract the key facts out of this text. Don't include opinions. Give each fact a number and keep them short sentences. :\n\n {text_input}"
#)

#def sentiment (  message):
#  sentiment_chain = LLMChain(llm=llm_hf_sentiment, prompt=sentiment_prompt)
#  facts = sentiment_chain.run(message)
#  print(facts)
#  return facts





####
## Chat models
# https://huggingface.co/spaces/CK42/sentiment-model-comparison
# 1 seem best for testing
####
chat_model_facebook_blenderbot_400M_distill = "facebook/blenderbot-400M-distill"
chat_model_HenryJJ_vincua_13b = "HenryJJ/vincua-13b"

text = "Why did the chicken cross the road?"

#output_question_1 = llm_hf(text)
#print(output_question_1)



###
## FACT EXTRACTION
###
# https://colab.research.google.com/drive/1hrS6_g14EcOD4ezwSGlGX2zxJegX5uNX#scrollTo=NUwUR9U7qkld
llm_factextract = HuggingFaceHub(
    
#    repo_id="google/flan-ul2", 
    repo_id="google/flan-t5-small", 
    model_kwargs={"temperature":0.1,
                  "max_new_tokens":250})
 
fact_extraction_prompt = PromptTemplate(
    input_variables=["text_input"],
    template="Extract the key facts out of this text. Don't include opinions. Give each fact a number and keep them short sentences. :\n\n {text_input}"
)

def factextraction (message):
  fact_extraction_chain = LLMChain(llm=llm_factextract, prompt=fact_extraction_prompt)
  facts = fact_extraction_chain.run(message)
  print(facts)
  return facts


####
##   models
# 1 seem best for testing
####
#download and setup the model and tokenizer
model_name = 'facebook/blenderbot-400M-distill'
tokenizer = BlenderbotTokenizer.from_pretrained(model_name)
model = BlenderbotForConditionalGeneration.from_pretrained(model_name)

def func (message):
  inputs = tokenizer(message, return_tensors="pt")
  result = model.generate(**inputs)
  print(result)
  return tokenizer.decode(result[0])



app = gr.Interface(
    fn=func,
    title="Conversation Bota",
    inputs=["text", "checkbox", gr.Slider(0, 100)],
    outputs=["text", "number"],
)
#app.launch()
####################


#app_sentiment = gr.Interface(fn=predict , inputs="textbox", outputs="textbox", title="Conversation Bot")
# create a public link, set `share=True` in `launch()
#app_sentiment.launch()
####################

###
###
###
classifier = pipeline("zero-shot-classification")
text = "This is a tutorial about Hugging Face."
candidate_labels = ["inform", "sell", "beschweren"]

def topic_sale_inform (text):
  res = classifier(text, candidate_labels)
  print (res)
  return res



####
conversation = Conversation("Welcome")

def callChains(current_message):
    sentiment_analysis_result = pipeline_predict_sentiment(current_message)
    topic_sale_inform_result = topic_sale_inform(current_message)
    #conversation.append_response("The Big lebowski.")
    #conversation.add_user_input("Is it good?")
    final_answer = func(current_message)
    return final_answer, sentiment_analysis_result, topic_sale_inform_result

chat_bot = gr.Interface(fn=callChains , inputs="textbox", outputs=["textbox","textbox","textbox"], title="Conversation Bot with extra")
# create a public link, set `share=True` in `launch()
chat_bot.launch()
####################


app_facts = gr.Interface(fn=factextraction , inputs="textbox", outputs="textbox", title="Conversation Bots")
# create a public link, set `share=True` in `launch()
#app_facts.launch()
####################