Upload folder using huggingface_hub

app.py

@@ -14,11 +14,12 @@ origins = {
                 'metier.format_court2']
 }
 
-def retrieve(origin='Formation', query='cuisine'):
+def retrieve(origin='Formation', query='cuisine', n_results=20):
     # Query API
     json = dict(
         query=query,
-        origins=origins[origin]
+        origins=origins[origin],
+        n_results=n_results
     )
 
     resp = requests.post(url=api_url('retrieve'), json=json)
@@ -33,11 +34,28 @@ def retrieve(origin='Formation', query='cuisine'):
     df['origin'] = df['origin'].apply(lambda x: x.split('.')[1])
     return df
 
+def rank(query='cuisine', documents = []):
+    # Query API
+    json = dict(
+        query=query,
+        documents=documents
+    )
+
+    try:
+        resp = requests.post(url=api_url('rank'), json=json)
+        data = resp.json()
+    except:
+        return []
+
+    # Format result
+    return data['ranking']
+
 gradio_app = gr.Interface(
     fn=retrieve,
     inputs=[
-        gr.Dropdown(list(origins.keys()), label="Origine", info="Choisir un type de donnée à interroger"),
-        gr.Textbox(label="Recherche", info="Votre recherche")
+        gr.Dropdown(choices=list(origins.keys()), value=list(origins.keys())[0], label="Origine", info="Choisir un type de donnée à interroger"),
+        gr.Textbox(label="Recherche", info="Votre recherche"),
+        gr.Number(value=10, label="Nombre de résultats", info="Nombre de résultats attendus")
     ],
     outputs=[
         gr.DataFrame(label="Résultats", headers=["Distance", "Key", "Label", "Origin", "Document"])

mps-api.py

@@ -43,10 +43,10 @@ class VECTORDB:
         print(f"{self.chroma_collection.count()} documents loaded.")
 
     @method()
-    def search(self, query, origins):
+    def search(self, query, origins, n_results=10):
         results = self.chroma_collection.query(
             query_texts=[query],
-            n_results=10,
+            n_results=n_results,
             where={"origin": {"$in": origins}},
             include=['documents', 'metadatas', 'distances'])
 
@@ -69,10 +69,8 @@ class RANKING:
     @method()
     def rank(self, query, documents):
         pairs = [[query, doc] for doc in documents]
-        print(pairs)
         scores = self.cross_encoder.predict(pairs)
-        print(scores)
-        ranking = np.argsort(scores)[::]
+        ranking = np.argsort(scores)[::].tolist()
         return ranking
 
 ###########
@@ -85,7 +83,7 @@ def retrieve(query: Dict):
     print(f"Retrieve query: {query}...")
 
     # Searching documents
-    documents, metadatas, distances = VECTORDB().search.remote(query['query'], query['origins'])
+    documents, metadatas, distances = VECTORDB().search.remote(query['query'], query['origins'], query['n_results'])
 
     return {"documents" : documents, "metadatas" : metadatas, "distances" : distances}