Upload folder using huggingface_hub

app.py

@@ -2,7 +2,8 @@ import gradio as gr
 import requests
 import pandas as pd
 
-api_url = 'https://huynhdoo--mps-api-query.modal.run'
+def api_url(remote):
+    return f"https://huynhdoo--mps-api-{remote}.modal.run"
 
 origins = {
     'Formation' : ['formation.presentation', 'formation.summary'],
@@ -13,14 +14,14 @@ origins = {
                 'metier.format_court2']
 }
 
-def API(origin='Formation', query='cuisine'):
+def retrieve(origin='Formation', query='cuisine'):
     # Query API
     json = dict(
         query=query,
         origins=origins[origin]
     )
 
-    resp = requests.post(url=api_url, json=json)
+    resp = requests.post(url=api_url('retrieve'), json=json)
     data = resp.json()
 
     # Format result
@@ -33,7 +34,7 @@ def API(origin='Formation', query='cuisine'):
     return df
 
 gradio_app = gr.Interface(
-    fn=API,
+    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")

mps-api.py

@@ -9,6 +9,7 @@ model_image = (Image.debian_slim(python_version="3.12")
 # Utilities
 with model_image.imports():
     import os
+    import numpy as np
     __import__("pysqlite3")
     sys.modules["sqlite3"] = sys.modules.pop("pysqlite3") # Hotswap SQLlite version
 
@@ -42,7 +43,7 @@ class VECTORDB:
         print(f"{self.chroma_collection.count()} documents loaded.")
 
     @method()
-    def query(self, query, origins):
+    def search(self, query, origins):
         results = self.chroma_collection.query(
             query_texts=[query],
             n_results=10,
@@ -54,18 +55,47 @@ class VECTORDB:
         distances = results['distances'][0]
         return documents, metadatas, distances
 
+@app.cls(timeout=30*60)
+class RANKING:
+    @enter()
+    @build()
+    def init(self):
+        # Load crossencoder
+        from sentence_transformers import CrossEncoder
+        model_name = "Lajavaness/CrossEncoder-camembert-large"
+        self.cross_encoder = CrossEncoder(model_name)
+        print(f"Cross encoder model loaded: {model_name}")
+    
+    @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)[::]
+        return ranking
+
 ###########
 # ENDPOINTS
 ###########
 @app.function(timeout=30*60)
 @web_endpoint(method="POST")
-def query(query: Dict):
+def retrieve(query: Dict):
     # Log query
-    print(f"Incoming query: {query}...")
+    print(f"Retrieve query: {query}...")
 
-    # Instantiate vectordb
-    vectordb = VECTORDB()
+    # Searching documents
+    documents, metadatas, distances = VECTORDB().search.remote(query['query'], query['origins'])
 
-    # Run query
-    documents, metadatas, distances = vectordb.query.remote(query['query'], query['origins'])
     return {"documents" : documents, "metadatas" : metadatas, "distances" : distances}
+
+@app.function(timeout=30*60)
+@web_endpoint(method="POST")
+def rank(query: Dict):
+    # Log query
+    print(f"Rank query: {query}...")
+
+    # Ranking documents
+    ranking = RANKING().rank.remote(query['query'], query['documents'])
+
+    return {"ranking" : ranking}