File size: 1,970 Bytes
8d89dbc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
import gradio as gr
import pandas as pd
import faiss
import numpy as np
import os
from sentence_transformers import SentenceTransformer
# Load the pre-trained embedding model
model = SentenceTransformer("Alibaba-NLP/gte-Qwen2-1.5B-instruct", trust_remote_code=True)
# Set the max sequence length
model.max_seq_length = 512
# Load the JSON data into a DataFrame
df = pd.read_json('White-Stride-Red-68.json')
df['embeding_context'] = df['embeding_context'].astype(str).fillna('')
# Filter out any rows where 'embeding_context' might be empty or invalid
df = df[df['embeding_context'] != '']
index = faiss.read_index('vector_store.index')
# Function to perform search and return all columns
def search_query(query_text):
num_records = 50
# Encode the input query text
embeddings_query = model.encode([query_text], prompt_name="query")
embeddings_query_np = np.array(embeddings_query).astype('float32')
# Search in FAISS index for nearest neighbors
distances, indices = index.search(embeddings_query_np, num_records)
# Get the top results based on FAISS indices
result_df = df.iloc[indices[0]].drop(columns=['embeding_context']).drop_duplicates().reset_index(drop=True)
return result_df
# Gradio interface function
def gradio_interface(query_text):
search_results = search_query(query_text)
return search_results
# Gradio interface setup
with gr.Blocks() as app:
gr.Markdown("<h1>White Stride Red Search (GTE-Qwen2)</h1>")
# Input text box for the search query
search_input = gr.Textbox(label="Search Query", placeholder="Enter search text", interactive=True)
# Output table for displaying results
search_output = gr.DataFrame(label="Search Results")
# Search button
search_button = gr.Button("Search")
# Link button click to action
search_button.click(fn=gradio_interface, inputs=search_input, outputs=search_output)
# Launch the Gradio app
app.launch()
|