Vectorize / app.py
0xalfroz's picture
Update app.py
56518e1 verified
raw
history blame
1.03 kB
import gradio as gr
from transformers import AutoModel, AutoTokenizer
import numpy as np
# Load a small CPU model for text to vector processing
model_name = "sentence-transformers/all-mpnet-base-v2"
model = AutoModel.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
def text_to_vector(texts):
# Tokenize and process a batch of inputs
inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True)
outputs = model(**inputs)
vectors = outputs.pooler_output.detach().numpy()
# Convert each vector to a string representation
vector_strings = [", ".join(map(str, vector)) for vector in vectors]
return vector_strings
demo = gr.Interface(
fn=text_to_vector,
inputs=gr.Textbox(label="Enter text", lines=4, placeholder="Enter one or more texts (separate with new lines)"),
outputs=gr.Textbox(label="Text Vectors", lines=10),
title="Batch Text to Vector",
description="This demo converts a batch of texts to vectors."
)
demo.launch()