Update app.py

app.py

@@ -1,80 +1,34 @@
-# app.py
-import os
-import json
 import pandas as pd
-import google.generativeai as genai
 import gradio as gr
-
-# Configure Gemini
-genai.configure(api_key=os.environ["GEMINI_API_KEY"])
-model = genai.GenerativeModel("gemini-2.0-flash")
+from retriever import get_relevant_passages
+from reranker import rerank
 
 # Load and clean CSV
-df_raw = pd.read_csv("assesments.csv")
-
-def preprocess_data(df):
-    def clean_duration(text):
-        try:
-            return int(text.split('=')[-1].strip())
-        except:
-            return None
-
-    def clean_support(val):
-        return "Yes" if val == 'T' else "No"
-
-    def clean_test_type(val):
-        return [x.strip() for x in str(val).split('\n') if x.strip()]
-
-    df_cleaned = pd.DataFrame({
-        "url": "https://www.shl.com" + df.iloc[:, 1].astype(str),
-        "remote_support": df.iloc[:, 2].apply(clean_support),
-        "adaptive_support": df.iloc[:, 3].apply(clean_support),
-        "test_type": df.iloc[:, 4].apply(clean_test_type),
-        "description": df.iloc[:, 5],
-        "duration": df.iloc[:, 8].apply(clean_duration),
-    })
-    return df_cleaned
-
-assessments = preprocess_data(df_raw)
-
-def recommend_assessments(query, top_k=10):
-    prompt = f"""
-    Given this job description: "{query}", recommend the top {top_k} relevant SHL assessments from the following list.
-    Return the result as JSON with this format:
-    {{
-        "recommended_assessments": [
-            {{
-                "url": ..., 
-                "adaptive_support": ..., 
-                "remote_support": ..., 
-                "description": ..., 
-                "duration": ..., 
-                "test_type": [...]
-            }},
-            ...
-        ]
-    }}
+df = pd.read_csv("assessments.csv")
 
-    Data:
-    {assessments.to_dict(orient='records')}
-    """
+def clean_df(df):
+    df = df.copy()
+    df["url"] = "https://www.shl.com" + df.iloc[:, 1]
+    df["remote_support"] = df.iloc[:, 2].map(lambda x: "Yes" if x == "T" else "No")
+    df["adaptive_support"] = df.iloc[:, 3].map(lambda x: "Yes" if x == "T" else "No")
+    df["test_type"] = df.iloc[:, 4].astype(str).str.split("\\n")
+    df["description"] = df.iloc[:, 5]
+    df["duration"] = df.iloc[:, 8].astype(str).str.extract(r'(\d+)').astype(float)
+    return df[["url", "adaptive_support", "remote_support", "description", "duration", "test_type"]]
 
-    response = model.generate_content(prompt)
-    try:
-        result = json.loads(response.text)
-        return result
-    except Exception as e:
-        return {"error": f"Failed to parse response: {str(e)}\n{response.text}"}
+df_clean = clean_df(df)
 
-def predict(query):
-    return recommend_assessments(query)
+def recommend(query):
+    top_k_df = get_relevant_passages(query, df_clean, top_k=20)
+    candidates = top_k_df.to_dict(orient="records")
+    return rerank(query, candidates)
 
 iface = gr.Interface(
-    fn=predict,
+    fn=recommend,
     inputs=gr.Textbox(label="Enter Job Description", lines=4),
     outputs="json",
-    title="SHL Assessment Recommender (Gemini-powered)",
-    description="Paste a job description and get the most relevant SHL assessments."
+    title="SHL Assessment Recommender",
+    description="Paste a job description to get the most relevant SHL assessments."
 )
 
 if __name__ == "__main__":