Create app.py

app.py

@@ -0,0 +1,114 @@
+import os
+import time
+import numpy as np
+import gradio as gr
+import torch
+import torch.nn.functional as F
+from transformers import AutoTokenizer, AutoModel
+
+MODEL = "intfloat/multilingual-e5-base"
+
+print("Loading tokenizer and model:", MODEL)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print("Device:", device)
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL)
+model = AutoModel.from_pretrained(MODEL)
+model.to(device)
+model.eval()
+
+# simple mean pooling using attention mask
+def mean_pooling(last_hidden_state, attention_mask):
+    input_mask_expanded = attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float()
+    sum_embeddings = torch.sum(last_hidden_state * input_mask_expanded, dim=1)
+    sum_mask = input_mask_expanded.sum(dim=1)
+    # avoid division by zero
+    sum_mask = torch.clamp(sum_mask, min=1e-9)
+    return sum_embeddings / sum_mask
+
+
+def embed_texts(texts, batch_size=8):
+    single = False
+    if isinstance(texts, str):
+        texts = [texts]
+        single = True
+
+    all_embs = []
+    with torch.no_grad():
+        for i in range(0, len(texts), batch_size):
+            batch = texts[i : i + batch_size]
+            encoded = tokenizer(batch, padding=True, truncation=True, max_length=512, return_tensors="pt")
+            input_ids = encoded["input_ids"].to(device)
+            attention_mask = encoded["attention_mask"].to(device)
+            outputs = model(input_ids=input_ids, attention_mask=attention_mask, return_dict=True)
+            last_hidden = outputs.last_hidden_state
+            pooled = mean_pooling(last_hidden, attention_mask)
+            normed = F.normalize(pooled, p=2, dim=1)
+            all_embs.append(normed.cpu().numpy())
+    embs = np.vstack(all_embs)
+    return embs[0] if single else embs
+
+
+def cosine_similarity(a, b):
+    # inputs are 1D arrays
+    denom = (np.linalg.norm(a) * np.linalg.norm(b))
+    if denom == 0:
+        return 0.0
+    return float(np.dot(a, b) / denom)
+
+
+def similarity(s1, s2):
+    e = embed_texts([s1, s2])
+    score = cosine_similarity(e[0], e[1])
+    return f"{score:.4f}"
+
+
+def search(query, docs_text, topk=3):
+    docs = [d.strip() for d in docs_text.splitlines() if d.strip()]
+    if not docs:
+        return "Corpus is empty"
+    all_texts = docs + [query]
+    embs = embed_texts(all_texts)
+    D, q = embs[:-1], embs[-1]
+    scores = (D @ q) / (np.linalg.norm(D, axis=1) * np.linalg.norm(q) + 1e-12)
+    order = np.argsort(scores)[::-1][:int(topk)]
+    lines = []
+    for rank, idx in enumerate(order, start=1):
+        lines.append(f"{rank}. score={scores[idx]:.4f}\n{docs[idx]}")
+    return "\n\n".join(lines)
+
+
+# Gradio UI
+
+demo = gr.Blocks()
+with demo:
+    gr.Markdown("# multilingual-e5-base — multilingual embedding tester")
+    with gr.Tab("Similarity"):
+        t1 = gr.Textbox(label="Text 1", value="Hello world / Привет мир")
+        t2 = gr.Textbox(label="Text 2", value="Greetings planet / Привет, планета")
+        btn = gr.Button("Compute similarity")
+        out = gr.Textbox(label="Cosine similarity", interactive=False)
+        btn.click(fn=similarity, inputs=[t1, t2], outputs=out)
+
+    with gr.Tab("Semantic search"):
+        q = gr.Textbox(label="Query", value="climate change")
+        corpus = gr.Textbox(label="Corpus (one document per line)", lines=12, value=(
+            "Climate summit discussed emissions reductions.\n"
+            "Local sports team won the championship.\n"
+            "New research on climate change effects published.\n"
+            "Economy grows despite challenges."))
+        k = gr.Number(label="Top-K", value=3, precision=0)
+        btn2 = gr.Button("Search")
+        out2 = gr.Textbox(label="Results", lines=12)
+        btn2.click(fn=search, inputs=[q, corpus, k], outputs=out2)
+
+    gr.Markdown("---\nModel: intfloat/multilingual-e5-base — uses Transformers AutoModel; runs on GPU if available.")
+
+# Launch
+if __name__ == "__main__":
+    # warmup: tokenize a small input to ensure weights moved to device
+    try:
+        _ = embed_texts(["Hello world"])  # warm cache
+    except Exception as e:
+        print("Warmup failed:", e)
+    demo.launch(server_name="0.0.0.0", server_port=7860)