Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from transformers import AutoModel, AutoTokenizer | |
| import torch | |
| model_name = "sentence-transformers/all-MiniLM-L6-v2" | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| model = AutoModel.from_pretrained(model_name) | |
| def get_embedding(text): | |
| inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True) | |
| with torch.no_grad(): | |
| outputs = model(**inputs) | |
| embeddings = outputs.last_hidden_state.mean(dim=1) | |
| return embeddings | |
| st.title("Text Embedding with all-MiniLM-L6-v2") | |
| st.write("Enter text to get its embedding:") | |
| input_text = st.text_input("Input Text", "") | |
| if input_text: | |
| embedding = get_embedding(input_text) | |
| st.write("Embedding:") | |
| st.write(embedding) | |