import gradio as gr

from load_model import load_Auto
from load_push import all_files
from retriever import *
from retrieve_docs import *
from make_chain_model import make_chain_llm
from make_answer import *
from transformers import TextStreamer

llm = load_Auto()
pinecone,bm25 = all_files('files')
retriever=retriever(pinecone,bm25)
rag_chain = make_chain_llm(retriever,llm)

def response(message, history):
    return rag_chain.invoke(message)

# def talk(,history):
#     # k = 1 # number of retrieved documents
#     # scores , retrieved_documents = search(prompt, k)
#     # formatted_prompt = format_prompt(prompt,retrieved_documents,k)
#     # formatted_prompt = formatted_prompt[:2000] # to avoid GPU OOM
    
#     messages =[
#     {"role": "system", "content": f"{PROMPT}"},
#     {}
#     {"role": "user", "content": f"{instruction}"}
#     ]
    
#     # tell the model to generate
#     input_ids = tokenizer.apply_chat_template(
#     messages,
#     add_generation_prompt=True,
#     return_tensors="pt"
#     ).to(model.device)

#     terminators = [
#     tokenizer.eos_token_id,
#     tokenizer.convert_tokens_to_ids("<|eot_id|>")
#     ]

#     text_streamer = TextStreamer(tokenizer)

#     outputs = model.generate(
#     input_ids,
#     max_new_tokens=4096,
#     eos_token_id=terminators,
#     do_sample=True,
#     streamer = text_streamer,
#     temperature=0.6,
#     top_p=0.9,
#     repetition_penalty = 1.1
#     )
#     response = []
#     for text in streamer:
#         response.append(outputs)
#         print(response)
#         yield "".join(response)
        

"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""

TITLE = "FUT FUT Chatbot"

DESCRIPTION = """
'해요'체를 사용하며 친절하게 답하는 풋풋이 챗봇.
A rag pipeline with a chatbot feature
Resources used to build this project :
* embedding model : https://huggingface.co/BM-K/KoSimCSE-roberta-multitask
* dataset : https://huggingface.co/datasets/Dongwookss/q_a_korean_futsal
* vector DB : PINECONE
* chatbot : https://huggingface.co/Dongwookss/small_fut_final
"""

Examples = [['시흥 풋살 구장 추천해줘'],['풋살 경기 규칙 설명해줘'], ['풋살 경기 시간 알려줘']]

demo = gr.ChatInterface(
    fn=response,
    chatbot=gr.Chatbot(
        show_label=True,
        show_share_button=True,
        show_copy_button=True,
        likeable=True,
        layout="bubble",
        bubble_full_width=False,
    ),
    theme="Soft",
    examples=Examples,
    title=TITLE,
    description=DESCRIPTION,
)
demo.launch(debug=True)

# if __name__ == "__main__":
#     demo.launch()