import gradio as gr
from transformers import pipeline
from sentence_transformers import SentenceTransformer
import spaces  # Import Hugging Face ZeroGPU module

# Load models
print("Loading BERT NER model...")
ner_pipeline = pipeline("ner", model="yashpwr/resume-ner-bert-v2", aggregation_strategy="simple")

print("Loading SBERT model...")
similarity_model = SentenceTransformer("all-MiniLM-L6-v2")

# Decorate GPU functions
@spaces.GPU
def parse_resume(text):
    if not text:
        return {"error": "Empty input"}
    try:
        results = ner_pipeline(text)
        formatted = []
        for entity in results:
            formatted.append({
                "entity_group": entity["entity_group"],
                "word": str(entity["word"]),
                "score": float(entity["score"])
            })
        return {"entities": formatted}
    except Exception as e:
        return {"error": str(e)}

@spaces.GPU
def get_embeddings(text):
    if not text:
        return {"error": "Empty input"}
    try:
        vector = similarity_model.encode(text).tolist()
        return {"embedding": vector}
    except Exception as e:
        return {"error": str(e)}

# Create Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# AI Resume Matcher Serverless API")
    
    with gr.Tab("NER Parser"):
        text_input = gr.Textbox(label="Resume Text", placeholder="Paste resume text here...", lines=10)
        parse_btn = gr.Button("Extract Entities")
        json_output = gr.JSON(label="Parsed JSON Output")
        parse_btn.click(fn=parse_resume, inputs=text_input, outputs=json_output, api_name="parse")
        
    with gr.Tab("Embeddings"):
        text_input_emb = gr.Textbox(label="Input Text", placeholder="Enter sentence or paragraph...", lines=5)
        embed_btn = gr.Button("Generate Vector")
        vector_output = gr.JSON(label="Embedding Vector")
        embed_btn.click(fn=get_embeddings, inputs=text_input_emb, outputs=vector_output, api_name="embed")

demo.launch()