path_work = "."

# hf_token
from dotenv import load_dotenv
load_dotenv()
import os
hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")


# [선택1] 거대모델 랭체인 Custom LLM (HF InferenceClient) - 70B가 무료!!!, openai보다 성능 안떨어짐 (스트리밍은 아직 안됨)
# model_name = "tiiuae/falcon-180B-chat"
model_name="meta-llama/Llama-2-70b-chat-hf"
# model_name="NousResearch/Llama-2-70b-chat-hf"
# model_name="meta-llama/Llama-2-13b-chat-hf"
# model_name="meta-llama/Llama-2-7b-chat-hf"
# model_name = "HuggingFaceH4/zephyr-7b-alpha"

kwargs = {"max_new_tokens":256, "temperature":0.9, "top_p":0.6, "repetition_penalty":1.3, "do_sample":True}

# 커스텀 LLM
from pydantic import BaseModel, Field
from typing import Any, Optional, Dict, List
from huggingface_hub import InferenceClient
from langchain.llms.base import LLM

class KwArgsModel(BaseModel):
    kwargs: Dict[str, Any] = Field(default_factory=dict)

class CustomInferenceClient(LLM, KwArgsModel):
    model_name: str
    inference_client: InferenceClient

    def __init__(self, model_name: str, hf_token: str, kwargs: Optional[Dict[str, Any]] = None):
        inference_client = InferenceClient(model=model_name, token=hf_token)
        super().__init__(
            model_name=model_name,
            hf_token=hf_token,
            kwargs=kwargs,
            inference_client=inference_client  # inference_client 인자 추가
        )

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None
    ) -> str:
        if stop is not None:
            raise ValueError("stop kwargs are not permitted.")
        # pdb.set_trace()
        # response_gen = self.__dict__['client'].text_generation(prompt, stream=True, **self.kwargs)  # 저장된 kwargs를 사용, 
        response_gen = self.inference_client.text_generation(prompt, **self.kwargs, stream=True)
        response = ''.join(response_gen)  # 제너레이터의 모든 값을 문자열로 연결
        return response

    @property
    def _llm_type(self) -> str:
        return "custom"

    @property
    def _identifying_params(self) -> dict:
        return {"model_name": self.model_name}

# 사용 예제:
# prompt="How do you make cheese?"
# prompt = "Tell me the names of the last 10 U.S. presidents"
prompt="Tell me 10 of the world's largest buildings in high order"

llm = CustomInferenceClient(model_name=model_name, hf_token=hf_token, kwargs=kwargs)  # hf_token 사용하는 경우
# llm = CustomInferenceClient(model_name=model_name, kwargs=kwargs)  # hf_token 사용하지 않는 경우


# 임베딩 객체 생성
from langchain.embeddings import HuggingFaceInstructEmbeddings
embeddings = HuggingFaceInstructEmbeddings(
    # model_name="sentence-transformers/all-MiniLM-L6-v2",  
    cache_folder="./sentence-transformers/all-MiniLM-L6-v2",
    model_kwargs={"device": "cpu"}
    )

# 벡터DB 로드
path_work ='.'

from langchain.vectorstores import Chroma
vectordb = Chroma(
    persist_directory = path_work + '/cromadb_llama2-papers',
    embedding_function=embeddings)

retriever = vectordb.as_retriever(search_kwargs={"k": 5}) 

# RetrievalQA 체인 만들기
from langchain.chains import RetrievalQA
qa_chain = RetrievalQA.from_chain_type(
    # llm=OpenAI(),  # from langchain.llms import OpenAI
    llm=llm,
    chain_type="stuff",
    retriever=retriever,
    return_source_documents=True,
    verbose=True,
    )    
qa_chain

# 그라디오
import json
import os
import gradio as gr

# Stream text
def predict(message, chatbot, temperature=0.9, max_new_tokens=512, top_p=0.6, repetition_penalty=1.3,):
        
    temperature = float(temperature)
    if temperature < 1e-2: temperature = 1e-2
    top_p = float(top_p)
    
    # 프롬프트
    # system_message = "\nYou are a psychological counselor who gives friendly and professional counseling on the concerns of Korean clients."
    # input_prompt = f"[INST] <<SYS>>\n{system_message}\n<</SYS>>\n\n "
    # for interaction in chatbot:
    #     input_prompt = input_prompt + str(interaction[0]) + " [/INST] " + str(interaction[1]) + " </s><s> [INST] "

    # input_prompt = input_prompt + str(message) + " [/INST] "


    # conversationalRetrievalChain (히스토리가 체인 내장 프롬프트에 인풋됨)
    # chat_history = []
    # for interaction in chatbot:
    #     chat_history = chat_history + [(str(interaction[0]), str(interaction[1]))]
    # llm_response = qa_chain_conv({"question": message, "chat_history": chat_history})
    # res_result = llm_response['answer']


    # RetrievalQA 체인 (히스토리가 체인 내장 프롬프트에 인풋 안됨)
    llm_response = qa_chain(message)
    res_result = llm_response['result']

    # conversationalRetrievalChain, RetrievalQA 체인 공통
    res_relevant_doc = [source.metadata['source'] for source in llm_response["source_documents"]]
    response = f"{res_result}" + "\n\n" + "[답변 근거 소스 논문 (ctrl + click 하세요!)] :" + "\n" + f" \n {res_relevant_doc}"
    print("response: =====> \n", response, "\n\n")

    #3) json 형태로 변환 (api response와 같은 형태)
    import json
    tokens = response.split('\n')
    token_list = []
    for idx, token in enumerate(tokens):
        token_dict = {"id": idx + 1, "text": token}
        token_list.append(token_dict)
    response = {"data": {"token": token_list}}
    response = json.dumps(response, indent=4)

    '''{'data': {'token': [{'id': 1, 'text': 'Artificial intelligence (AI) refers to...'},
                           {'id': 2, 'text': 'I hope this information helher questions!'}]}}'''

    # ===========================================================================
    # 스트리밍 시작 (partial_message)
    response = json.loads(response)  # {'data': {'token': [{'id': 1, 'text': '답변은 " 안녕하세요. 저는 송상진 박사.....
    data_dict = response.get('data', {})
    token_list = data_dict.get('token', [])

    import time
    partial_message = ""
    # 하이라이트: .iter_lines() 대신에 token_list를 직접 순회합니다.
    for token_entry in token_list:
        if token_entry:  # filter out keep-alive new lines (if any)
            try:
                # 하이라이트: 직접 사전에서 'id'와 'text'를 추출합니다.
                token_id = token_entry.get('id', None)
                token_text = token_entry.get('text', None)

                # time.sleep으로 글자 속도 조절하며 글자 내보냄
                if token_text:  # 이 부분은 원하는 대로 조정할 수 있습니다.
                    # partial_message = partial_message + token_text
                    for char in token_text:  # 문자 하나씩 순회 (추가됨)
                        partial_message += char  # partial_message에 문자 추가 (변경됨)
                        yield partial_message
                        time.sleep(0.01)
                else:
                    # gr.Warning(f"The key 'text' does not exist or is None in this token entry: {token_entry}")
                    print(f"[[워닝]] ==> The key 'text' does not exist or is None in this token entry: {token_entry}")

            except KeyError as e:
                gr.Warning(f"KeyError: {e} occurred for token entry: {token_entry}")
                continue

# 타이틀/설명/질문예시
title = "llama-2 모델 관련 논문 QA 서비스"
description = """chat history 유지 보다는 QA에 충실하도록 제작되었으니 Single turn으로 활용을 하여 주세요. (chat history 활용은 다른 주제로 별도 제작 예정)"""
css = """.toast-wrap { display: none !important } """
examples=[['Can you tell me about the llama-2 model?'],['What is percent accuracy, using the SPP layer as features on the SPP (ZF-5) model?'], ['What is percent accuracy, using the SPP layer as features on the SPP (ZF-5) model?'], ["tell me about method for human pose estimation based on DNNs"]]

# 좋아요
import gradio as gr
def vote(data: gr.LikeData):
    if data.liked: print("You upvoted this response: " + data.value)
    else: print("You downvoted this response: " + data.value)

# 그라디오 (인자 조절)
additional_inputs = [
    # gr.Textbox("", label="Optional system prompt"),
    gr.Slider(label="Temperature", value=0.9, minimum=0.0, maximum=1.0, step=0.05, interactive=True, info="Higher values produce more diverse outputs"),
    gr.Slider(label="Max new tokens", value=256, minimum=0, maximum=4096, step=64, interactive=True, info="The maximum numbers of new tokens"),
    gr.Slider(label="Top-p (nucleus sampling)", value=0.6, minimum=0.0, maximum=1, step=0.05, interactive=True, info="Higher values sample more low-probability tokens"),
    gr.Slider(label="Repetition penalty", value=1.2, minimum=1.0, maximum=2.0, step=0.05, interactive=True, info="Penalize repeated tokens")
]

chatbot_stream = gr.Chatbot(avatar_images=(
    "https://drive.google.com/uc?id=13rYrN0cH_9tR7GveqO1q2JiyBCqkfCLZ", # https://drive.google.com/uc?id= 뒤에 ID값만 (모두 사용자 액세스 권한 허용)
    "https://drive.google.com/uc?id=1tfELAQW_VbPCy6QTRbexRlwAEYo8rSSv"
    ), bubble_full_width = False)

chat_interface_stream = gr.ChatInterface(predict, 
                 title=title, 
                 description=description, 
                #  textbox=gr.Textbox(lines=5),
                 chatbot=chatbot_stream,
                 css=css, 
                 examples=examples, 
                #  cache_examples=True, 
                #  additional_inputs=additional_inputs,
                 ) 

# Gradio Demo 
with gr.Blocks() as demo:

    with gr.Tab("스트리밍"):
        #gr.ChatInterface(predict, title=title, description=description, css=css, examples=examples, cache_examples=True, additional_inputs=additional_inputs,) 
        chatbot_stream.like(vote, None, None)
        chat_interface_stream.render()

        
demo.queue(concurrency_count=75, max_size=100).launch(debug=True)