from sentence_transformers import SentenceTransformer
import gradio as gr
import faiss
import numpy as np
model = SentenceTransformer('sentence-transformers/all-mpnet-base-v2')
#model.save('models/all-mpnet-base-v2')
#model.to('cuda')
def EmdeddingVect(sentences):
    return model.encode([sentences])  # Assume model.encode() function exists

# Load the arrays
all_embeddings = np.load('all_embeddings.npy')
all_labels = np.load('all_labels.npy')

index = faiss.IndexFlatL2(768)
index.add(all_embeddings.astype('float32'))

def search_query(k, query_sentence):
    query_embeddings = EmdeddingVect(query_sentence)
    distances, indices = index.search(query_embeddings.astype('float32'), int(k))
    ai_count = np.sum(all_labels[indices[0]] == 'AI')
    ai_probability = (ai_count / int(k)) * 100
    human_probability = 100 - ai_probability
    return f"Probability of being AI: {ai_probability:.2f}%", f"Probability of being Human: {human_probability:.2f}%"

iface = gr.Interface(
    fn=search_query, 
    inputs=[
        gr.inputs.Slider(minimum=1, maximum=10, default=3, label="k (Number of Neighbors)"),
        gr.inputs.Textbox(label="Text to be Detected")
    ], 
    outputs=[
        gr.outputs.Textbox(label="AI Probability"),
        gr.outputs.Textbox(label="Human Probability")
    ]
)

iface.launch()