Upload 2 files

app.py

@@ -0,0 +1,201 @@
+from pymilvus import MilvusClient, connections
+from pymilvus import Collection, DataType, FieldSchema, CollectionSchema, utility
+
+from sentence_transformers import SentenceTransformer
+import re
+
+import gradio as gr
+import openai
+
+import requests
+import json
+
+def connect_milvus(endpoint, token):
+    connections.connect(
+        "default",
+        uri=CLUSTER_ENDPOINT,
+        token=TOKEN
+    )
+    
+    client = MilvusClient(
+        uri=CLUSTER_ENDPOINT,
+        token=TOKEN
+    )
+    #print(client)
+    
+def create_vectordb(COLLECTION_NAME, EMBEDDING_DIM=1024):
+    
+    COLLECTION_NAME = "van_ban_phap_luat"
+    check_collection = utility.has_collection(COLLECTION_NAME)
+    if check_collection:
+        drop_result = utility.drop_collection(COLLECTION_NAME)
+        print('drop completed!')
+        
+    chunk_id = FieldSchema(
+        name="chunk_id",
+        dtype=DataType.INT64,
+        is_primary=True,
+        description="Chunk ID"
+    )
+
+    chunk_ref = FieldSchema(
+        name="chunk_ref",
+        dtype=DataType.VARCHAR,
+        max_length=512,
+        description="Chunk ref")
+
+    chunk_text = FieldSchema(
+        name="chunk_text",
+        dtype=DataType.VARCHAR,
+        max_length=4096,
+        description="Chunk text")
+
+    chunk_embedding = FieldSchema(
+        name="chunk_embedding",
+        dtype=DataType.FLOAT_VECTOR,
+        dim=EMBEDDING_DIM,
+        description="Chunk Embedding")
+
+    schema = CollectionSchema(
+        fields=[chunk_id, chunk_ref, chunk_text, chunk_embedding],
+        auto_id=False,
+        description="Vector Store Chunk using multilingual-e5-large")
+
+    collection = Collection(
+        name=COLLECTION_NAME,
+        schema=schema
+    )
+    
+    entities = [ids, rules, chunks, embeddings]
+    ins_resp = collection.insert(entities)
+    ins_resp
+    
+    collection.flush()
+    
+    index_params = {
+        "index_type": "IVF_FLAT",
+        "metric_type": "COSINE",  # L2
+        "params": {}
+    }
+    collection.create_index(
+        field_name=chunk_embedding.name,
+        index_params=index_params
+    )
+    
+    collection.load()
+    #return None   
+
+def load_vectordb(COLLECTION_NAME):
+    
+    collection = Collection(
+        name=COLLECTION_NAME,
+        #schema=schema
+    )
+    #print(collection)
+    #print(collection.has_index())
+    return collection
+
+def load_model(model_name):
+    model = SentenceTransformer(model_name)
+
+    return model
+
+def search_chunks(query, topk=5):
+        # search
+    search_params = {
+        "metric_type": "COSINE",  # L2
+        "params": {"level": 2}
+    }
+
+    collection = load_vectordb(COLLECTION_NAME)
+
+    embed_query = model.encode(query)
+    results = collection.search(
+        [embed_query],
+        #anns_field=chunk_embedding.name,
+        anns_field="chunk_embedding",
+        param=search_params,
+        limit=topk,
+        guarantee_timestamp=1,
+        output_fields=['chunk_ref', 'chunk_text'] #
+    )
+
+    refs, relevant_chunks = [], []
+
+    pattern = r"'chunk_ref': '([^']*)', 'chunk_text': '([^']*)'"
+
+    #print(len(results[0]))
+
+    for sample in results[0]:
+        #print(sample)
+        #print(type(sample))
+        matches = re.findall(pattern, str(sample))
+
+        # Lặp qua các kết quả tìm được và thêm vào list
+        for match in matches:
+            refs.append(match[0])
+            relevant_chunks.append(match[1])
+
+    return refs, relevant_chunks
+
+    #return results[0]
+    #print(len(results[0]))
+    #return results
+
+CLUSTER_ENDPOINT = "https://in03-d63abe0e8a8f47b.api.gcp-us-west1.zillizcloud.com"
+TOKEN = "6529d1f59d5e3d38d6135ec9ddf5820a9c38e3db6ca22c53b3aa2ad9c9148e29ef0f11c312bee71f1544da350f85320b598a30f3"
+COLLECTION_NAME = "van_ban_phap_luat"
+MODEL_NAME = 'qminh369/datn-dense_embedding'
+#query = "Cơ quan quản lý giáo dục và đào tạo phải làm gì để đảm bảo pháp luật về giao thông đường bộ được đưa vào chương trình giảng dạy?" # 7
+#query = "Trách nhiệm và quản lý hoạt động giao thông đường bộ được phân công và phân cấp như thế nào?" # 4
+#query = "Những trường hợp cho phép chở người trên xe ô tô chở hàng?" # 21
+
+connect_milvus(CLUSTER_ENDPOINT, TOKEN)
+#load_vectordb(COLLECTION_NAME)
+model = load_model(MODEL_NAME)
+
+model_gpt_3_5 = "gpt-3.5-turbo"
+openai.api_key = "sk-proj-P6SydZ6cYprqmpMH6BqdT3BlbkFJT4hjCdGZpGZPfcpBqi7q"
+
+def write_prompt(doc_relevant, question):
+    
+    PROMPT = "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n{question}\n\n### Input:\n{doc_relevant}"
+    
+    input_prompt = PROMPT.format_map(  
+    {"doc_relevant": doc_relevant, "question": question}  
+)
+    return input_prompt
+
+def answer(question):
+
+    refs, relevant_chunks = search_chunks(question)
+    #print(results)
+    system_msg = "Hãy trả lời câu hỏi sau dựa trên thông tin được cung cấp. Nếu thông tin được cung cấp không liên quan dến câu hỏi thì trả về câu trả lời 'Không có thông tin'. Lưu ý các câu trả lời đều sử dụng Tiếng Việt"
+    user_msg = write_prompt(relevant_chunks, question)
+    
+    response = openai.ChatCompletion.create(model=model_gpt_3_5,
+                                            messages=[{"role": "system", "content": system_msg},
+                                             {"role": "user", "content": user_msg}])
+    
+    response = response["choices"][0]["message"]["content"]
+    
+    ref = "\n" + "Trích dẫn từ: " + refs[0]
+    
+    response = response + ref
+    
+    return response.strip()
+
+def chatbot(question, history=[]):
+    output = answer(question)
+    history.append((question, output))
+    return history, history
+
+demo = gr.Interface(
+    fn=chatbot,
+    inputs=["text", "state"],
+    outputs=["chatbot", "state"])
+
+demo.queue().launch(share=True)
+
+
+

requirements.txt

@@ -0,0 +1,5 @@
+sentence-transformers
+openai
+grpcio==1.60.0
+pymilvus
+gradio