import gradio as gr
import torch
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer
import spaces 

MODEL_NAME = "Qwen/Qwen3-Embedding-0.6B" 

print(f"Loading {MODEL_NAME} to RAM...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModel.from_pretrained(MODEL_NAME)
model.eval()
print("Model ready!")

@spaces.GPU
def get_embedding(text):
    # [PERBAIKAN]: Paksa model pindah ke GPU saat fungsi ini dieksekusi oleh ZeroGPU
    model.to("cuda")
    
    # 1. Tokenisasi dan pindahkan input ke GPU
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
    inputs = {k: v.to("cuda") for k, v in inputs.items()}
    
    # 2. Forward pass
    with torch.no_grad():
        outputs = model(**inputs)
    
    # 3. Mean Pooling dengan Attention Mask
    attention_mask = inputs['attention_mask']
    last_hidden_state = outputs.last_hidden_state
    
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float()
    sum_embeddings = torch.sum(last_hidden_state * input_mask_expanded, 1)
    sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
    embeddings = sum_embeddings / sum_mask
    
    # 4. L2 Normalization
    embeddings = F.normalize(embeddings, p=2, dim=1)
    
    return embeddings.squeeze().tolist()

demo = gr.Interface(
    fn=get_embedding,
    inputs=gr.Textbox(lines=3, placeholder="Masukkan teks untuk di-embed..."),
    outputs="json",
    title="Qwen3 Embedding 0.6B API (ZeroGPU)"
)

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