trhhh

app.py

@@ -1,21 +1,18 @@
 import os
-import requests
 import gradio as gr
 import chromadb
-import json
-import pandas as pd
 
 from sentence_transformers import SentenceTransformer
 
 import spaces
 
 @spaces.GPU
-def get_embeddings(text, task):
+def get_embeddings(query, task):
     model = SentenceTransformer("Linq-AI-Research/Linq-Embed-Mistral", use_auth_token=os.getenv("HF_TOKEN"))
     task = "Given a question, retrieve passages that answer the question"
-    prompt = f"Instruct: {task}\nQuery: {text}"  # Use text here
-    query_embeddings = model.encode([prompt], convert_to_tensor=True)  # Ensure it's a list
-    return query_embeddings.cpu().numpy()
+    prompt = f"Instruct: {task}\nQuery: {query}" 
+    query_embeddings = model.encode([prompt])  
+    return query_embeddings
 
 
 # Initialize a persistent Chroma client and retrieve collection
@@ -56,27 +53,28 @@ def query_chroma(embeddings, authors, num_results=10):
 
         return formatted_results
     except Exception as e:
-        return f"Failed to query the database: {str(e)}"
+        return {"error": str(e)}
 
 
 # Main function
-def perform_query(query, task, author, num_results):
+def perform_query(query, authors, num_results):
+    task = "Given a question, retrieve passages that answer the question"
     embeddings = get_embeddings(query, task)
-    initial_results = query_chroma(embeddings, author, num_results)
+    results = query_chroma(embeddings, authors, num_results)
+
+    if "error" in results:
+        return [gr.update(visible=True, value=f"Error: {results['error']}") for _ in range(max_textboxes * 3)]
 
-    results = [(f"{res['author']}, {res['book']}, Distance: {res['distance']}", res['text'], res['id']) for res in initial_results]
-    
     updates = []
-    for meta, text, id_ in results:
-        markdown_content = f"**{meta}**\n\n{text}"
+    for res in results:
+        markdown_content = f"**{res['author']}, {res['book']}, Distance: {res['distance']}**\n\n{res['text']}"
         updates.append(gr.update(visible=True, value=markdown_content))
         updates.append(gr.update(visible=True, value="Flag", elem_id=f"flag-{len(updates)//2}"))
-        updates.append(gr.update(visible=False, value=id_))  # Hide the ID textbox
+        updates.append(gr.update(visible=False, value=res['id']))  # Hide the ID textbox
     
-    updates += [gr.update(visible=False)] * (3 * (max_textboxes - len(results) // 3))
-    
-    return updates
+    updates += [gr.update(visible=False)] * (3 * (max_textboxes - len(results)))
 
+    return updates
 
 # Initialize the CSVLogger callback for flagging 
 callback = gr.CSVLogger()