Update app.py

app.py

@@ -1,11 +1,30 @@
+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()
 
+path_work = "."
 hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
 
+embeddings = HuggingFaceInstructEmbeddings(
+    model_name="sentence-transformers/all-MiniLM-L6-v2",  
+    model_kwargs={"device": "cpu"}
+    )
+
+vectordb = Chroma(
+    persist_directory = path_work + '/cromadb_llama2-papers',
+    embedding_function=embeddings)
+
+retriever = vectordb.as_retriever(search_kwargs={"k": 5})
 
 class KwArgsModel(BaseModel):
     kwargs: Dict[str, Any] = Field(default_factory=dict)
@@ -42,4 +61,124 @@ class CustomInferenceClient(LLM, KwArgsModel):
     def _identifying_params(self) -> dict:
         return {"model_name": self.model_name}
 
-kwargs = {"max_new_tokens":256, "temperature":0.9, "top_p":0.6, "repetition_penalty":1.3, "do_sample":True}
+kwargs = {"max_new_tokens":256, "temperature":0.9, "top_p":0.6, "repetition_penalty":1.3, "do_sample":True}
+
+model_list=[
+    "meta-llama/Llama-2-13b-chat-hf",   
+    "HuggingFaceH4/zephyr-7b-alpha",
+    "meta-llama/Llama-2-70b-chat-hf",    
+    "tiiuae/falcon-180B-chat"
+]
+
+qa_chain = None
+
+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} 로딩 완료."
+
+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)
+    
+    llm_response = qa_chain(message)
+    res_result = llm_response['result']
+
+    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")
+
+    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)
+
+    response = json.loads(response)  
+    data_dict = response.get('data', {})
+    token_list = data_dict.get('token', [])
+
+    partial_message = ""
+    for token_entry in token_list:
+        if token_entry:  
+            try:
+                token_id = token_entry.get('id', None)
+                token_text = token_entry.get('text', None)
+
+                if token_text:  
+                    for char in token_text:  
+                        partial_message += char  
+                        yield partial_message
+                        time.sleep(0.01)
+                else:
+                    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 모델 관련 논문 Generative 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?'], ["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"]]
+
+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.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=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, 
+                chatbot=chatbot_stream,
+                css=css, 
+                examples=examples, 
+                 ) 
+
+with gr.Blocks() as demo:
+    with gr.Tab("스트리밍"):
+        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)