Upload folder using huggingface_hub

app.py

@@ -0,0 +1,128 @@
+import json
+import argparse
+from pathlib import Path
+from typing import List
+
+import gradio as gr
+import faiss
+import numpy as np
+import torch
+from sentence_transformers import SentenceTransformer
+
+
+file_example = """Please upload a JSON file with a "text" field (with optional "title" field). For example
+```JSON
+[
+    {"title": "", "text": "This an example text without the title"},
+    {"title": "Title A", "text": "This an example text with the title"},
+    {"title": "Title B", "text": "This an example text with the title"},
+]
+```"""
+
+
+def create_index(embeddings, use_gpu):
+    index = faiss.IndexFlatIP(len(embeddings[0]))
+    embeddings = np.asarray(embeddings, dtype=np.float32)
+    if use_gpu:
+        co = faiss.GpuMultipleClonerOptions()
+        co.shard = True
+        co.useFloat16 = True
+        index = faiss.index_cpu_to_all_gpus(index, co=co)
+    index.add(embeddings)
+    return index
+
+
+def upload_file_fn(
+        file_path: List[str], 
+        progress: gr.Progress = gr.Progress(track_tqdm=True)
+    ):
+    try:
+        with open(file_path) as f:
+            document_data = json.load(f)
+        documents = []
+        for obj in document_data:
+            text = obj["title"] + "\n" + obj["text"] if obj.get("title") else obj["text"]
+            documents.append(text)
+    except Exception as e:
+        print(e)
+        gr.Warning("Read the file failed. Please check the data format.")
+        return None, None
+    
+    documents_embeddings = model.encode(documents)
+
+    document_index = create_index(documents_embeddings, use_gpu=False)
+    
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+    
+    return document_index, document_data
+
+
+def clear_file_fn():
+    return None, None
+
+
+def retrieve_document_fn(question, document_data, document_index):
+    num_retrieval_doc = 3
+    if document_index is None or document_data is None:
+        gr.Warning("Please upload documents first!")
+        return [None for i in range(num_retrieval_doc)]
+    
+    question_embedding = model.encode([question])
+    batch_scores, batch_inxs = document_index.search(question_embedding, k=num_retrieval_doc)
+    
+    answers = [document_data[i]["text"] for i in batch_inxs[0][:num_retrieval_doc]]
+    return tuple(answers)
+
+
+def main(args):
+    global model
+
+    model = SentenceTransformer(args.model_name_or_path)
+
+    document_index = gr.State()
+    document_data = gr.State()
+
+    
+    with open(Path(__file__).parent / "resources/head.html") as html_file:
+        head = html_file.read().strip()
+    with gr.Blocks(theme=gr.themes.Soft(font="sans-serif").set(background_fill_primary="linear-gradient(90deg, #e3ffe7 0%, #d9e7ff 100%)", background_fill_primary_dark="linear-gradient(90deg, #4b6cb7 0%, #182848 100%)",),
+                   head=head, 
+                   css=Path(__file__).parent / "resources/styles.css", 
+                   title="KaLM-Embedding", 
+                   fill_height=True, 
+                   analytics_enabled=False) as demo:
+        gr.Markdown(file_example)
+        doc_files_box = gr.File(label="Upload Documents", file_types=[".json"], file_count="single")
+        retrieval_interface = gr.Interface(
+            fn=retrieve_document_fn,
+            inputs=["text"],
+            outputs=["text", "text",  "text"],
+            additional_inputs=[document_data, document_index],
+            concurrency_limit=1,
+        )
+
+        doc_files_box.upload(
+            upload_file_fn, 
+            [doc_files_box], 
+            [document_index, document_data],
+            queue=True,
+            trigger_mode="once"
+        )
+        doc_files_box.clear(
+                upload_file_fn, 
+                None, 
+                [document_index, document_data],
+                queue=True,
+                trigger_mode="once"
+            )
+    demo.launch()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_name_or_path", type=str, default="HIT-TMG/KaLM-embedding-multilingual-mini-instruct-v1.5")
+
+    args = parser.parse_args()
+    main(args)

requirements.txt

@@ -0,0 +1,3 @@
+transformers==4.39.1
+sentence-transformers==2.5.1
+faiss-cpu==1.8.0

resources/head.html

@@ -0,0 +1,36 @@
+<style>
+:root {
+--highlight-background-color-light: #F1EAFF;
+--highlight-background-color-dark: #3E00FF;
+}
+
+body {
+--highlight-background-color: var(--highlight-background-color-light);
+}
+
+@media (prefers-color-scheme: dark) {
+body {
+  --highlight-background-color: var(--highlight-background-color-dark);
+}
+}
+</style>
+
+<script>
+document.addEventListener('click', function(event) {
+  if (event.target.tagName.toLowerCase() === 'a') {
+  var href = event.target.getAttribute('href');
+
+  if (href && href.startsWith('#')) {
+      var targetId = href.substring(1);
+      var targetArticle = document.getElementById(targetId);
+      var articles = document.getElementsByTagName('article');
+      
+      for (var i = 0; i < articles.length; i++) {
+          articles[i].style.backgroundColor = '';
+      }
+      
+      targetArticle.style.backgroundColor = 'var(--highlight-background-color)';
+  }
+}
+});
+</script>

resources/styles.css

@@ -0,0 +1,158 @@
+
+a {
+  text-decoration: none!important;
+}
+
+article {
+  padding: 12px;
+}
+
+
+progress::-webkit-progress-bar {
+  border-radius: 8rpx !important;
+  background-color: #f0f0f0;
+}
+progress::-webkit-progress-value {
+  border-radius: 8rpx !important;
+}
+.progress-1 {
+  color: #FF004D; 
+}
+.progress-1::-webkit-progress-value {
+  background-color:#FF004D;
+}
+.progress-1::-moz-progress-bar {
+  background-color:#FF004D;
+}
+.progress-2 {
+  color: #FF8400; 
+}
+.progress-2::-webkit-progress-value {
+  background-color:#FF8400;
+}
+.progress-2::-moz-progress-bar {
+  background-color:#FF8400;
+}
+.progress-3 {
+  color: #0079FF; 
+}
+.progress-3::-webkit-progress-value {
+  background-color:#0079FF;
+}
+.progress-3::-moz-progress-bar {
+  background-color:#0079FF;
+}
+
+.factual-score {
+  width: 20px;
+  height: 4px;
+}
+
+
+.hide {
+  display: none;
+}
+
+.html-text {
+  white-space: pre-line;
+  word-break: normal;
+  text-align: justify;
+  overflow-y: auto;
+  overflow-x: hidden;
+  height: 450px;
+  padding: 2px;
+}
+
+.tab {
+  border: none;
+  outline: none;
+  padding: 0;
+  margin: 0;
+}
+
+
+/* html text slider */
+::-webkit-scrollbar {
+  width: 6px;
+  height: 6px;
+}
+::-webkit-scrollbar-track {
+  border-radius: 3px;
+  background: rgba(0,0,0,0.06);
+  -webkit-box-shadow: inset 0 0 5px rgba(0,0,0,0.08);
+}
+::-webkit-scrollbar-thumb {
+  border-radius: 3px;
+  background: rgba(0,0,0,0.12);
+  -webkit-box-shadow: inset 0 0 10px rgba(0,0,0,0.2);
+}
+
+
+/* split line style */
+.hr-edge-weak {
+  border: 0;
+  padding-top: 1px;
+  background: linear-gradient(to right, transparent, #d0d0d5, transparent);
+}
+
+.hr-double-arrow {
+  color: #d0d0d5;
+  border: double;
+  border-width: 3px 5px;
+  border-color: #d0d0d5 transparent;
+  height: 1px;
+  overflow: visible;
+  margin-left: 20px;
+  margin-right: 20px;
+  position: relative;
+}
+.hr-double-arrow:before, 
+.hr-double-arrow:after {
+  content: '';
+  position: absolute;
+  width: 5px; height: 5px;
+  border-width: 0 3px 3px 0;
+  border-style: double;
+  top: -3px;
+  background: radial-gradient(2px at 1px 1px, currentColor 2px, transparent 0) no-repeat;
+}
+.hr-double-arrow:before {
+  transform: rotate(-45deg);
+  left: -20px;
+}
+.hr-double-arrow:after {
+  transform: rotate(135deg);
+  right: -20px;
+}
+
+
+.citation-button {
+  -webkit-tap-highlight-color: rgba(0,0,0,0);
+  -webkit-text-size-adjust: 100%;
+  tab-size: 4;
+  color-scheme: light;
+  word-break: break-word;
+  white-space: pre-wrap;
+  font-family: Open Sans,sans-serif!important;
+  box-sizing: border-box;
+  border: 0 solid #e5e7eb;
+  touch-action: manipulation;
+  margin-right: 1px;
+  display: inline-flex;
+  height: .75rem;
+  min-width: .75rem;
+  align-items: center;
+  justify-content: center;
+  border-radius: 9999px;
+  background-color: var(--block-label-background-fill);
+  padding-left: .25rem;
+  padding-right: .25rem;
+  vertical-align: top;
+  font-size: 8px;
+  font-weight: 600;
+  line-height: 0;
+  color: var(--block-label-text-color);
+  outline: none;
+  cursor: pointer;
+  transition: color 0.3s;
+}