ThesisLast / src /streamlit_app.py
NotRev's picture
Update src/streamlit_app.py
04e4783 verified
import streamlit as st
from transformers import pipeline
import json
import os
# --- Page Configuration ---
st.set_page_config(page_title="Skill vs Knowledge Extractor", layout="wide")
@st.cache_resource
def load_models():
# Load NER (Finds the terms) and Zero-Shot Classifier (Categorizes them)
try:
st.info("Loading AI Models (Hugging Face local models)... This may take a moment.")
# Model 1: Named Entity Recognition for finding candidate terms
# CORRECTED MODEL ID: "jjzha/jobbert-base-cased"
ner_pipe = pipeline("token-classification",
model="jjzha/jobbert-base-cased",
aggregation_strategy="simple")
# Model 2: Zero-Shot Classification for categorizing terms
classifier_pipe = pipeline("zero-shot-classification",
model="valhalla/distilbart-mnli-12-1")
return ner_pipe, classifier_pipe
except Exception as e:
# Note: If the error persists, check your internet connection and ensure
# your device has enough memory to download these large models.
st.error(f"FATAL: Error loading models. Ensure 'transformers', 'accelerate', 'streamlit', and 'torch' are installed. Details: {e}")
return None, None
def process_text(text, ner_pipe, classifier_pipe):
if not text:
return {"SKILL": [], "KNOWLEDGE": []}
# 1. Extract Candidates (Using NER Model)
ner_results = ner_pipe(text)
candidates = set()
for entity in ner_results:
word = entity['word'].strip()
# Filter out short or single-character entities
if len(word.split()) > 1 or len(word) > 2:
candidates.add(word)
candidates = list(candidates)
if not candidates:
return {"SKILL": [], "KNOWLEDGE": []}
# --- THESIS ENHANCEMENT: Heuristic Post-Processing Overrides ---
# These lists are used to correct the known (and often variable) biases
# of the zero-shot classifier for specific technical terms.
SKILL_OVERRIDES = ["RAG", "function calling", "LoRA", "CI/CD pipelines", "DeepEval", "RAGAS", "Azure", "AWS"]
KNOWLEDGE_OVERRIDES = ["clean code practices", "English fluency", "async code", "team leadership", "agile methodologies"]
skills, knowledge = [], []
classification_labels = ["software tool or technology", "concept or knowledge"]
for candidate in candidates:
# Check Overrides First (Highest priority for accuracy)
if candidate in SKILL_OVERRIDES:
skills.append(candidate)
continue
if candidate in KNOWLEDGE_OVERRIDES:
knowledge.append(candidate)
continue
# 2. Classify (Zero-Shot Model)
try:
result = classifier_pipe(candidate, candidate_labels=classification_labels)
top_label = result['labels'][0]
# The zero-shot model determines the category
if top_label == "software tool or technology":
skills.append(candidate)
else:
knowledge.append(candidate)
except Exception as e:
# Fallback for errors or empty results
knowledge.append(candidate)
return {
"SKILL": sorted(list(set(skills))),
"KNOWLEDGE": sorted(list(set(knowledge)))
}
# --- UI Layout ---
st.title("💡 AI Job Description Analyzer")
ner_pipe, classifier_pipe = load_models()
if ner_pipe and classifier_pipe:
st.markdown("""
***Methodology:*** *This application uses a two-stage NLP pipeline: 1) The `jjzha/jobbert-base-cased` NER model to identify relevant terms, followed by 2) The `valhalla/distilbart-mnli-12-1` Zero-Shot Classifier to categorize them as 'SKILL' or 'KNOWLEDGE'. A heuristic post-processing layer ensures high precision for key technical terms.*
""")
job_description = st.text_area(
"Job Description Text",
height=300,
placeholder="Paste a job description here..."
)
if st.button("Analyze and Extract Entities", type="primary"):
if job_description.strip():
with st.spinner("Analyzing text and running classification..."):
output = process_text(job_description, ner_pipe, classifier_pipe)
st.subheader("Extraction Output (JSON)")
st.json(output)
json_str = json.dumps(output, indent=2)
st.download_button(
label="Download JSON Output",
data=json_str,
file_name="extracted_entities.json",
mime="application/json"
)
else:
st.warning("Please paste a job description into the text area.")