app.py

import gradio as gr
from transformers import AutoTokenizer, AutoModel
import torch

# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-Embedding-0.6B")
model = AutoModel.from_pretrained("Qwen/Qwen3-Embedding-0.6B")

def get_embedding(text):
    if len(text) > 250:
        return "❌ Error: Input exceeds 250 character limit."

    inputs = tokenizer(text, return_tensors="pt", truncation=True)
    with torch.no_grad():
        outputs = model(**inputs)
        # Use [CLS] token embedding (or mean pooling)
        embedding = outputs.last_hidden_state[:, 0, :].squeeze().tolist()

    # Show only first 10 dimensions for readability
    return f"✅ Embedding (first 10 values): {embedding[:10]}..."

demo = gr.Interface(
    fn=get_embedding,
    inputs=gr.Textbox(label="Enter a sentence (max 250 characters)", max_lines=3, placeholder="Type your sentence here...", lines=2),
    outputs="text",
    title="Qwen3 Embedding Demo",
    description="Generates sentence embeddings using Qwen/Qwen3-Embedding-0.6B. Input must be 250 characters or fewer."
)

if __name__ == "__main__":
    demo.launch()