Create app.py

app.py

@@ -0,0 +1,59 @@
+import time
+import os
+import gradio as gr
+import torch
+from transformers import AutoModel, AutoTokenizer
+import meilisearch
+
+tokenizer = AutoTokenizer.from_pretrained('Snowflake/snowflake-arctic-embed-m')
+model = AutoModel.from_pretrained('Snowflake/snowflake-arctic-embed-m', add_pooling_layer=False)
+model.eval()
+
+cuda_available = torch.cuda.is_available()
+print(f"CUDA available: {cuda_available}")
+
+meilisearch_client = meilisearch.Client("https://edge.meilisearch.com", os.environ["MEILISEARCH_KEY"])
+meilisearch_index_name = "docs-embed"
+meilisearch_index = meilisearch_client.index(meilisearch_index_name)
+
+def search_embeddings(query_text):
+    start_time_embedding = time.time()
+    query_prefix = 'Represent this sentence for searching code documentation: '
+    query_tokens =  tokenizer(query_prefix + query_text, padding=True, truncation=True, return_tensors='pt', max_length=512)
+    # step1: tokenizer the query
+    with torch.no_grad():
+        # Compute token embeddings
+        query_embeddings = model(**query_tokens)[0][:, 0]
+        # normalize embeddings
+        query_embeddings = torch.nn.functional.normalize(query_embeddings, p=2, dim=1)
+        document_embeddings_list = query_embeddings[0].tolist()
+        elapsed_time_embedding = time.time() - start_time_embedding
+    
+    # step2: search meilisearch
+    start_time_meilisearch = time.time()
+    response = meilisearch_index.search(
+        "", opt_params={"vector": document_embeddings_list, "hybrid": {"semanticRatio": 1.0}, "limit": 5, "attributesToRetrieve": ["text", "source", "library"]}
+    )
+    elapsed_time_meilisearch = time.time() - start_time_meilisearch
+    hits = response["hits"]
+
+    # step3: present the results in markdown
+    md = f"Stats:\n\nembedding time: {elapsed_time_embedding:.2f}s\n\nmeilisearch time: {elapsed_time_meilisearch:.2f}s\n\n---\n\n"
+    for hit in hits:
+        text, source, library = hit["text"], hit["source"], hit["library"]
+        source = f"[source](https://huggingface.co/docs/{library}/{source})"
+        md += text + f"\n\n{source}\n\n---\n\n"
+
+    return md
+
+
+demo = gr.Interface(
+    fn=search_embeddings,
+    inputs=gr.Textbox(label="enter your query", placeholder="Type Markdown here...", lines=10),
+    outputs=gr.Markdown(),
+    title="HF Docs Emebddings Explorer",
+    allow_flagging="never"
+)
+
+if __name__ == "__main__":
+    demo.launch()