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

MODEL_NAME = "microsoft/codebert-base"

# 初始化 tokenizer 和 model
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModel.from_pretrained(MODEL_NAME)

# 建立一個 feature-extraction pipeline
feature_extractor = pipeline(
    task="feature-extraction",
    model=model,
    tokenizer=tokenizer,
    return_tensors=True,   # 讓輸出保持 tensor 形式
    framework="pt"        # 指定 PyTorch
)

def embed_text(input_text):
    # pipeline 預設輸出的是 List[List[List[float]]] 的三維陣列
    # 可依需求做 pooling
    outputs = feature_extractor(input_text)
    # 假設只取第一個 batch 並做平均池化
    # outputs[0] => List[List[float]]
    # 其中每個 token 都是一個向量
    token_embeddings = torch.tensor(outputs[0])
    embedding = torch.mean(token_embeddings, dim=0)  # 取所有 token 的平均
    return embedding.tolist()

# 用 Gradio 建立一個介面
with gr.Blocks():
    gr.Markdown("# Feature Extraction Demo with CodeBERT")
    text_input = gr.Textbox(label="Enter text here")
    output = gr.JSON(label="Vector Embedding")

    btn = gr.Button("Compute Embedding")
    btn.click(embed_text, inputs=text_input, outputs=output)

    gr.Examples(
        examples=["def hello_world(): print('Hello World!')", "print('test')"],
        inputs=text_input,
        outputs=output
    )

    gr.launch()