Update app.py

app.py

@@ -1,60 +1,59 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import os
 import gradio as gr
-import PyPDF2
 import docx2txt
+import PyPDF2
+from keybert.llm import OpenAI
+from keybert import KeyLLM
 
-# Load token and model
-hf_token = os.environ.get("hf_space_token")
-model_name = "google/gemma-3-1b-it"
-tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
-model = AutoModelForCausalLM.from_pretrained(model_name, token=hf_token)
+# Get OpenAI key from environment
+api_key = os.environ.get("OPENAI_API_KEY")
 
-# Resume text extraction
-def extract_text_from_resume(file):
+# Setup OpenAI-backed KeyLLM model
+llm = OpenAI(api_key=api_key)
+kw_model = KeyLLM(llm)
+
+# Function to extract text from uploaded resume
+def extract_text(file):
     if file is None:
-        return "No file uploaded."
+        return ""
     ext = file.name.split('.')[-1].lower()
-    if ext == "pdf":
+    if ext == 'pdf':
         reader = PyPDF2.PdfReader(file)
         return "\n".join([page.extract_text() for page in reader.pages if page.extract_text()])
-    elif ext == "docx":
+    elif ext == 'docx':
         return docx2txt.process(file.name)
     else:
         return "Unsupported file format."
 
-# Function to generate LLM response
-def analyze_resume(job_description, resume_file, analyze_with_jd):
-    resume_text = extract_text_from_resume(resume_file)
+# Function to extract keywords using KeyLLM
+def extract_keywords_from_resume(job_desc, resume_file, analyze_with_jd):
+    resume_text = extract_text(resume_file)
     if not resume_text.strip():
-        return "Unable to extract resume content."
-
-    if analyze_with_jd and job_description.strip():
-        prompt = (
-            f"Compare the following resume with this job description:\n\n"
-            f"Job Description:\n{job_description}\n\n"
-            f"Resume:\n{resume_text}\n\n"
-            f"Give a short summary of how well this resume matches the job."
-        )
-    else:
-        prompt = f"Summarize the following resume:\n\n{resume_text}"
+        return "Could not extract text from resume."
 
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048)
-    outputs = model.generate(**inputs, max_new_tokens=200, do_sample=True, temperature=0.7)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # Create combined document if analysis is requested
+    document = job_desc + "\n\n" + resume_text if analyze_with_jd and job_desc else resume_text
 
-# Build Gradio UI
+    # Use KeyLLM to extract keywords
+    keywords = kw_model.extract_keywords([document], check_vocab=True)
+    return ", ".join(keywords[0]) if keywords else "No keywords found."
+
+# Gradio UI
 with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             analyze_checkbox = gr.Checkbox(label="Analyze with Job Description", value=True)
-            job_desc = gr.Textbox(label="Job Description", lines=8, placeholder="Paste job description here...")
+            job_desc = gr.Textbox(label="Job Description", lines=6)
             resume_file = gr.File(label="Upload Resume (PDF or DOCX)", file_types=[".pdf", ".docx"])
         with gr.Column():
-            parsed_output = gr.Textbox(label="Gemma Analysis Output", lines=20)
+            output_keywords = gr.Textbox(label="Extracted Keywords", lines=10)
 
-    resume_file.change(fn=analyze_resume, inputs=[job_desc, resume_file, analyze_checkbox], outputs=parsed_output)
-    job_desc.change(fn=analyze_resume, inputs=[job_desc, resume_file, analyze_checkbox], outputs=parsed_output)
+    resume_file.change(fn=extract_keywords_from_resume,
+                       inputs=[job_desc, resume_file, analyze_checkbox],
+                       outputs=output_keywords)
 
-demo.launch()
+    job_desc.change(fn=extract_keywords_from_resume,
+                    inputs=[job_desc, resume_file, analyze_checkbox],
+                    outputs=output_keywords)
 
+demo.launch()