Update app.py

app.py

@@ -1,36 +1,18 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModel
-import torch
+from sentence_transformers import SentenceTransformer
 
 print("Loading IndoBERT model...")
 MODEL_NAME = "indobenchmark/indobert-base-p1"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModel.from_pretrained(MODEL_NAME)
-model.eval()
+model = SentenceTransformer(MODEL_NAME)
 print("Model loaded!")
 
-def mean_pooling(model_output, attention_mask):
-    token_embeddings = model_output[0]
-    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
-    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
-
-def generate_embedding(text):
-    encoded_input = tokenizer(text, padding=True, truncation=True, max_length=512, return_tensors='pt')
-    
-    with torch.no_grad():
-        model_output = model(**encoded_input)
-    
-    embedding = mean_pooling(model_output, encoded_input['attention_mask'])
-    embedding = torch.nn.functional.normalize(embedding, p=2, dim=1)
-    
-    return embedding[0].numpy().tolist()
-
 def embed_single(text):
     """For Gradio interface - single text"""
     if not text:
         return {"error": "Text required"}
     
-    embedding = generate_embedding(text)
+    embedding = model.encode(text, normalize_embeddings=True).tolist()
+    
     return {
         "success": True,
         "embedding": embedding,
@@ -43,7 +25,7 @@ def embed_batch(texts):
         return {"error": "Texts required"}
     
     text_list = [t.strip() for t in texts.split('\n') if t.strip()]
-    embeddings = [generate_embedding(text) for text in text_list]
+    embeddings = model.encode(text_list, normalize_embeddings=True).tolist()
     
     return {
         "success": True,
@@ -57,14 +39,18 @@ with gr.Blocks() as demo:
     gr.Markdown("# 🇮🇩 IndoBERT Embedding API")
     
     with gr.Tab("Single"):
-        input_single = gr.Textbox(label="Text", lines=3)
-        btn_single = gr.Button("Generate")
+        input_single = gr.Textbox(label="Text", lines=3, placeholder="Enter Indonesian text...")
+        btn_single = gr.Button("Generate Embedding")
         output_single = gr.JSON(label="Result")
         btn_single.click(embed_single, inputs=input_single, outputs=output_single)
     
     with gr.Tab("Batch"):
-        input_batch = gr.Textbox(label="Texts (one per line)", lines=10)
-        btn_batch = gr.Button("Generate Batch")
+        input_batch = gr.Textbox(
+            label="Texts (one per line)", 
+            lines=10,
+            placeholder="Enter multiple Indonesian texts, one per line..."
+        )
+        btn_batch = gr.Button("Generate Batch Embeddings")
         output_batch = gr.JSON(label="Result")
         btn_batch.click(embed_batch, inputs=input_batch, outputs=output_batch)