Add sentence-transformers to requirements

main.py

@@ -1,38 +0,0 @@
-# This is a sample Python script.
-
-# Press Shift+F10 to execute it or replace it with your code.
-# Press Double Shift to search everywhere for classes, files, tool windows, actions, and settings.
-
-import gradio as gr
-import os
-
-from reranker.reranker import CrossEncReranker
-from retriever.es_retriever import ESRetriever
-from utils.preprocessing import question_to_statement
-
-
-ES_HOST = os.environ["ES_HOST"]
-ES_INDEX_NAME = os.environ["ES_INDEX_NAME"]
-ES_USERNAME = os.environ["ES_USERNAME"]
-ES_PASSWORD = os.environ["ES_PASSWORD"]
-
-RERANKER_MODEL_NAME = "douglasfaisal/granularity-legal-reranker-cross-encoder-indobert-base-p2"
-
-es_retriever_client = ESRetriever(ES_HOST, ES_INDEX_NAME, ES_USERNAME, ES_PASSWORD)
-cross_enc_reranker = CrossEncReranker(RERANKER_MODEL_NAME, 512)
-
-def retrieve_and_rerank(question: str):
-
-    query = question_to_statement(question)
-    retrieval_results = es_retriever_client.retrieve(query)
-    reranker_results = cross_enc_reranker.rerank(query, retrieval_results)
-
-    return reranker_results[0].text
-
-
-demo = gr.Interface(fn=retrieve_and_rerank, inputs="text", outputs="text")
-
-# Press the green button in the gutter to run the script.
-demo.launch()
-
-# See PyCharm help at https://www.jetbrains.com/help/pycharm/

requirements.txt

@@ -1,3 +1,4 @@
 gradio~=3.28.1
 numpy==1.21.4
 requests==2.26.0
+sentence-transformers

reranker/reranker.py

@@ -10,7 +10,7 @@ class CrossEncReranker:
         self.reranker = CrossEncoder(self.model_name)
         self.reranker.max_length = max_length
 
-    def rerank(self, query_text: str, candidates: list[LawComponent]):
+    def rerank(self, query_text: str, candidates: list):
         sentence_combinations = [[query_text, c.text] for c in candidates]
         similarity_scores = self.reranker.predict(sentence_combinations)
         index = np.argsort(similarity_scores)[::-1]