qwen-embed-app / app.py
anji29's picture
Create app.py
6428a42 verified
raw
history blame contribute delete
702 Bytes
import gradio as gr
from transformers import AutoTokenizer, AutoModel
import torch
model_name = "Qwen/Qwen-Embedding-0.6B"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
def get_embedding(text):
inputs = tokenizer(text, return_tensors="pt")
with torch.no_grad():
embeddings = model(**inputs).last_hidden_state.mean(dim=1).squeeze().tolist()
return embeddings
demo = gr.Interface(fn=get_embedding,
inputs=gr.Textbox(lines=2, placeholder="Enter text..."),
outputs="json",
title="Qwen Embedding Generator")
demo.launch()