app.py

import streamlit as st
from sentence_transformers import SentenceTransformer
import chromadb
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import time
from PyPDF2 import PdfReader

st.set_page_config(page_title="HireMate", page_icon="🧠")
st.title("🧠 HireMate - GenAI Talent Screener")
st.caption("Open-source RAG-based resume screener built without Azure or paid APIs")

# Initialize embedding model + vector DB
@st.cache_resource
def load_embedder():
    return SentenceTransformer("all-MiniLM-L6-v2")

@st.cache_resource
def load_llm():
    model_name = "mistralai/Mistral-7B-Instruct-v0.2"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16)
    return tokenizer, model

embedder = load_embedder()
tokenizer, model = load_llm()
client = chromadb.Client()
collection = client.create_collection("hiremate_resumes")

# Input JD
jd_text = st.text_area("📋 Paste Job Description", height=200)

# Upload resumes
uploaded_files = st.file_uploader("📄 Upload Resumes (PDF or Text)", accept_multiple_files=True)

def extract_text(file):
    if file.name.endswith(".pdf"):
        reader = PdfReader(file)
        return "".join([p.extract_text() for p in reader.pages])
    else:
        return file.read().decode("utf-8")

def score_resume(jd, resume):
    prompt = f"""
You are a hiring assistant. Evaluate how well this resume matches the job description.
Give a score (1–10) and one-line justification.

Job Description:
{jd}

Resume:
{resume}

Respond in the format:
Score: <number>/10 - <justification>
"""
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, max_new_tokens=120)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

if st.button("🚀 Screen Candidates"):
    if jd_text and uploaded_files:
        start = time.time()

        # Create JD embedding
        jd_emb = embedder.encode([jd_text]).tolist()
        collection.add(documents=[jd_text], embeddings=jd_emb, ids=["JD_1"])

        # Add resumes to Chroma
        for file in uploaded_files:
            text = extract_text(file)
            emb = embedder.encode([text]).tolist()
            collection.add(documents=[text], embeddings=emb, ids=[file.name])

        # Retrieve top matches
        results = collection.query(query_embeddings=jd_emb, n_results=min(5, len(uploaded_files)))

        st.subheader("🔍 Top Matching Resumes")
        for i, doc in enumerate(results["documents"][0]):
            with st.expander(f"Candidate {i+1} – Preview"):
                st.write(doc[:400] + "...")
                st.write("⏳ Scoring with LLM…")
                response = score_resume(jd_text, doc)
                st.success(response)

        latency = round(time.time() - start, 2)
        st.info(f"✅ Screening completed in {latency} seconds")
    else:
        st.warning("Please provide a Job Description and at least one resume.")