Update app.py

app.py

@@ -51,6 +51,11 @@ def clear_cuda_cache():
 def free_memory(*args):
     for arg in args:
         del arg
+
+def load_corpus_from_json(file_path):
+    with open(file_path, 'r') as file:
+        data = json.load(file)
+    return data
     
 @spaces.GPU
 def compute_embeddings(selected_task, input_text):
@@ -100,10 +105,8 @@ def compute_similarity(selected_task, sentence1, sentence2, extra_sentence1, ext
     free_memory(embeddings1, embeddings2, embeddings3, embeddings4)
 
     similarity_scores = {"Similarity 1-2": similarity1, "Similarity 1-3": similarity2, "Similarity 1-4": similarity3}
-
     return similarity_scores
-
-
+    
 @spaces.GPU
 def compute_cosine_similarity(emb1, emb2):
     tensor1 = torch.tensor(emb1).to(device).half()
@@ -112,7 +115,36 @@ def compute_cosine_similarity(emb1, emb2):
     free_memory(tensor1, tensor2)
     return similarity
 
+
+@spaces.GPU
+def compute_embeddings_batch(input_texts):
+    max_length = 2042
+    processed_texts = [f'Instruct: {task_description}\nQuery: {text}' for text in input_texts]
+    
+    batch_dict = tokenizer(processed_texts, max_length=max_length - 1, return_attention_mask=False, padding=False, truncation=True)
+    batch_dict['input_ids'] = [input_ids + [tokenizer.eos_token_id] for input_ids in batch_dict['input_ids']]
+    batch_dict = tokenizer.pad(batch_dict, padding=True, return_attention_mask=True, return_tensors='pt')
+    batch_dict = {k: v.to(device) for k, v in batch_dict.items()}
+    outputs = model(**batch_dict)
+    embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
+    embeddings = F.normalize(embeddings, p=2, dim=1)
+    return embeddings.detach().cpu().numpy()
+
+def semantic_search(query_embedding, corpus_embeddings, top_k=5):
+    scores = np.dot(corpus_embeddings, query_embedding.T).flatten()
+    top_k_indices = np.argsort(scores)[::-1][:top_k]
+    return top_k_indices, scores[top_k_indices]
+
+def search_similar_sentences(input_question, corpus_sentences, corpus_embeddings):
+    question_embedding = compute_embeddings_batch([input_question])[0]
+    top_k_indices, top_k_scores = semantic_search(question_embedding, corpus_embeddings)
+    results = [(corpus_sentences[i], top_k_scores[i]) for i in top_k_indices]
+    return results
+
+
 def app_interface():
+    corpus_sentences = []
+    corpus_embeddings = []
     with gr.Blocks() as demo:
         gr.Markdown(title)
         gr.Markdown(description)
@@ -141,6 +173,43 @@ def app_interface():
                 inputs=[task_dropdown, sentence1_box, sentence2_box, extra_sentence1_box, extra_sentence2_box],
                 outputs=similarity_output
             )
+        with gr.Tab("Load Corpus"):
+            json_uploader = gr.File(label="Upload JSON File")
+            load_corpus_button = gr.Button("Load Corpus")
+            corpus_status = gr.Textbox(label="Corpus Status", value="Corpus not loaded")
+
+            def load_corpus(file_info):
+                if file_info is None:
+                    return "No file uploaded. Please upload a JSON file."
+                try:
+                    global corpus_sentences, corpus_embeddings
+                    corpus_sentences = load_corpus_from_json(file_info['name'])
+                    corpus_embeddings = compute_embeddings_batch(corpus_sentences)
+                    return "Corpus loaded successfully with {} sentences.".format(len(corpus_sentences))
+                except Exception as e:
+                    return "Error loading corpus: {}".format(e)
+
+            load_corpus_button.click(
+                fn=load_corpus,
+                inputs=json_uploader,
+                outputs=corpus_status
+            )
+
+        with gr.Tab("Semantic Search"):
+            input_question_box = gr.Textbox(label="Enter your question")
+            search_button = gr.Button("Search")
+            search_results_output = gr.Textbox(label="Search Results")
+
+            def perform_search(input_question):
+                if not corpus_sentences or not corpus_embeddings:
+                    return "Corpus is not loaded. Please load a corpus first."
+                return search_similar_sentences(input_question, corpus_sentences, corpus_embeddings)
+
+            search_button.click(
+                fn=perform_search,
+                inputs=input_question_box,
+                outputs=search_results_output
+            )
 
         with gr.Row():
             with gr.Column():