Update app.py

app.py

@@ -1,27 +1,30 @@
-path_work = "."
-
-# hf_token
+import json
+import os
+import gradio as gr
+import time
+from pydantic import BaseModel, Field
+from typing import Any, Optional, Dict, List
+from huggingface_hub import InferenceClient
+from langchain.llms.base import LLM
+from langchain.embeddings import HuggingFaceInstructEmbeddings
+from langchain.vectorstores import Chroma
+import os
 from dotenv import load_dotenv
 load_dotenv()
-import os
-hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
-# os.environ['HUGGINGFACE_TOKEN'] = os.getenv("HUGGINGFACEHUB_API_TOKEN")
 
+path_work = "."
+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="meta-llama/Llama-2-13b-chat-hf"
-# model_name="meta-llama/Llama-2-7b-chat-hf"
-# model_name = "HuggingFaceH4/zephyr-7b-alpha"
+embeddings = HuggingFaceInstructEmbeddings(
+    model_name="sentence-transformers/all-MiniLM-L6-v2",  
+    model_kwargs={"device": "cpu"}
+    )
 
-kwargs = {"max_new_tokens":256, "temperature":0.9, "top_p":0.6, "repetition_penalty":1.3, "do_sample":True}
+vectordb = Chroma(
+    persist_directory = path_work + '/cromadb_llama2-papers',
+    embedding_function=embeddings)
 
-# 커스텀 LLM
-from pydantic import BaseModel, Field
-from typing import Any, Optional, Dict, List
-from huggingface_hub import InferenceClient
-from langchain.llms.base import LLM
+retriever = vectordb.as_retriever(search_kwargs={"k": 5})
 
 class KwArgsModel(BaseModel):
     kwargs: Dict[str, Any] = Field(default_factory=dict)
@@ -31,13 +34,12 @@ class CustomInferenceClient(LLM, KwArgsModel):
     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)
         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 인자 추가
+            inference_client=inference_client
         )
 
     def _call(
@@ -47,10 +49,8 @@ class CustomInferenceClient(LLM, KwArgsModel):
     ) -> 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)  # 제너레이터의 모든 값을 문자열로 연결
+        response = ''.join(response_gen)  
         return response
 
     @property
@@ -61,85 +61,51 @@ class CustomInferenceClient(LLM, KwArgsModel):
     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"}
-    )
+kwargs = {"max_new_tokens":256, "temperature":0.9, "top_p":0.6, "repetition_penalty":1.3, "do_sample":True}
 
-# 벡터DB 로드
-path_work ='.'
+model_list=[
+    "meta-llama/Llama-2-13b-chat-hf",   
+    "HuggingFaceH4/zephyr-7b-alpha",
+    "meta-llama/Llama-2-70b-chat-hf",    
+    "tiiuae/falcon-180B-chat"
+]
 
-from langchain.vectorstores import Chroma
-vectordb = Chroma(
-    persist_directory = path_work + '/cromadb_llama2-papers',
-    embedding_function=embeddings)
+qa_chain = None
 
-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
+def load_model(model_selected):
+    global qa_chain
+    model_name = model_selected
+    llm = CustomInferenceClient(model_name=model_name, hf_token=hf_token, kwargs=kwargs) 
+    
+    from langchain.chains import RetrievalQA
+    qa_chain = RetrievalQA.from_chain_type(
+        llm=llm,
+        chain_type="stuff",
+        retriever=retriever,
+        return_source_documents=True,
+        verbose=True,
+        )    
+    qa_chain
+
+load_model("meta-llama/Llama-2-70b-chat-hf")
+
+def model_select(model_selected):
+    load_model(model_selected)
+    return f"모델 {model_selected} 로딩 완료."
 
-# 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):
@@ -148,55 +114,40 @@ def predict(message, chatbot, temperature=0.9, max_new_tokens=512, top_p=0.6, re
     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': '답변은 " 안녕하세요. 저는 송상진 박사.....
+    response = json.loads(response)  
     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)
+        if token_entry:  
             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에 문자 추가 (변경됨)
+                if token_text:  
+                    for char in token_text:  
+                        partial_message += char  
                         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}")
+                    pass
 
             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 활용은 다른 주제로 별도 제작 예정)"""
+title = "Llama-2 모델 관련 논문 Generaatie QA (with RAG) 서비스 (Llama-2 70b 모델 활용)"
+description = """Chat history 유지 보다는 QA에 충실하도록 제작되었으므로 Single turn으로 활용 하여 주세요. Default로 Llama-2 70b 모델로 설정되어 있으나 GPU 서비스 한도 초과로 Error가 발생할 수 있으니 양해부탁드리���, 화면 하단의 모델 변경/로딩하시어 다른 모델로 변경하여 사용을 부탁드립니다. (다만, Llama-2 70b가 가장 정확하오니 참고하여 주시기 바랍니다.) """
 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"]]
+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?'], ["How much less accurate is using the SPP layer as features on the SPP (ZF-5) model compared to using the same model on the undistorted full image?"], ["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"),
@@ -204,29 +155,30 @@ additional_inputs = [
 ]
 
 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=18xKoNOHN15H_qmGhK__VKnGjKjirrquW",
     "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,
+chat_interface_stream = gr.ChatInterface(
+                predict,
+                title=title, 
+                description=description, 
+                chatbot=chatbot_stream,
+                css=css, 
+                examples=examples, 
                  ) 
 
-# 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()
-
+        with gr.Row():
+            with gr.Column(scale=6):
+                with gr.Row():
+                    model_selector = gr.Dropdown(model_list, label="모델 선택", value= "meta-llama/Llama-2-70b-chat-hf", scale=5)
+                    submit_btn1 = gr.Button(value="모델 로드", scale=1)
+            with gr.Column(scale=4):
+                model_status = gr.Textbox(value="", label="모델 상태")
+                submit_btn1.click(model_select, inputs=[model_selector], outputs=[model_status])
         
-demo.queue(concurrency_count=75, max_size=100).launch(debug=True)
-
+demo.queue(concurrency_count=75, max_size=100).launch(debug=True)