Update retriever.py

retriever.py

@@ -1,38 +1,110 @@
 import pandas as pd
-from sentence_transformers import SentenceTransformer, util
+import gradio as gr
+from retriever import get_relevant_passages
+from reranker import rerank
 
-model = SentenceTransformer("all-MiniLM-L6-v2")
-
-def get_relevant_passages(query, df, top_k=20):
-    # Create a copy to avoid modifying the original dataframe
-    df_copy = df.copy()
+# Load and clean CSV
+def clean_df(df):
+    df = df.copy()
+    
+    # Ensure clean URLs
+    # Check if the second column contains URLs or just IDs
+    second_col = df.iloc[:, 1].astype(str)
+    if second_col.str.contains('http').any() or second_col.str.contains('www').any():
+        df["url"] = second_col  # Already has full URLs
+    else:
+        # Create full URLs from IDs
+        df["url"] = "https://www.shl.com/" + second_col.str.replace(r'^[\s/]*', '', regex=True)
     
-    # Ensure URL field is properly formatted
-    if 'url' in df_copy.columns:
-        # Clean up URLs if needed
-        df_copy['url'] = df_copy['url'].astype(str)
-        # Ensure URLs start with http or https
-        mask = ~df_copy['url'].str.startswith(('http://', 'https://'))
-        df_copy.loc[mask, 'url'] = 'https://www.shl.com/' + df_copy.loc[mask, 'url'].str.lstrip('/')
+    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")
     
-    # Format test_type for better representation
-    def format_test_type(test_types):
-        if isinstance(test_types, list):
-            return ', '.join(test_types)
-        return str(test_types)
+    # Handle test_type with error checking
+    df["test_type"] = df.iloc[:, 4].astype(str).str.split("\\n")
     
-    # Concatenate all fields into a single string per row
-    corpus = df_copy.apply(
-        lambda row: f"{row['description']} "
-                   f"Test types: {format_test_type(row['test_type'])}. "
-                   f"Adaptive support: {row['adaptive_support']}. "
-                   f"Remote support: {row['remote_support']}. "
-                   f"Duration: {row['duration'] if pd.notna(row['duration']) else 'N/A'} minutes.",
-        axis=1
-    ).tolist()
+    df["description"] = df.iloc[:, 5]
     
-    corpus_embeddings = model.encode(corpus, convert_to_tensor=True)
-    query_embedding = model.encode(query, convert_to_tensor=True)
-    hits = util.semantic_search(query_embedding, corpus_embeddings, top_k=top_k)[0]
+    # Extract duration with error handling
+    df["duration"] = pd.to_numeric(
+        df.iloc[:, 8].astype(str).str.extract(r'(\d+)')[0], 
+        errors='coerce'
+    )
+    
+    return df[["url", "adaptive_support", "remote_support", "description", "duration", "test_type"]]
+
+try:
+    df = pd.read_csv("assesments.csv")
+    df_clean = clean_df(df)
+except Exception as e:
+    print(f"Error loading or cleaning data: {e}")
+    # Create an empty DataFrame with required columns as fallback
+    df_clean = pd.DataFrame(columns=["url", "adaptive_support", "remote_support", 
+                                     "description", "duration", "test_type"])
+
+def validate_and_fix_urls(candidates):
+    """Validates and fixes URLs in candidate assessments."""
+    for candidate in candidates:
+        # Ensure URL exists
+        if 'url' not in candidate or not candidate['url']:
+            candidate['url'] = 'https://www.shl.com/missing-url'
+            continue
+            
+        url = str(candidate['url'])
+        
+        # Fix URLs that are just numbers
+        if url.isdigit() or (url.startswith('https://www.shl.com') and url[len('https://www.shl.com'):].isdigit()):
+            candidate['url'] = f"https://www.shl.com/{url.replace('https://www.shl.com', '')}"
+            continue
+            
+        # Add protocol if missing
+        if not url.startswith(('http://', 'https://')):
+            candidate['url'] = f"https://{url}"
+            
+    return candidates
+
+def recommend(query):
+    if not query.strip():
+        return {"error": "Please enter a job description"}
     
-    return df_copy.iloc[[hit['corpus_id'] for hit in hits]]
+    try:
+        # Print some debug info
+        print(f"Processing query: {query[:50]}...")
+        
+        top_k_df = get_relevant_passages(query, df_clean, top_k=20)
+        
+        # Debug: Check URLs in retrieved data
+        print(f"Retrieved {len(top_k_df)} assessments")
+        if not top_k_df.empty:
+            print(f"Sample URLs from retrieval: {top_k_df['url'].iloc[:3].tolist()}")
+        
+        candidates = top_k_df.to_dict(orient="records")
+        
+        # Additional URL validation before sending to reranker
+        for c in candidates:
+            if 'url' in c:
+                if not str(c['url']).startswith(('http://', 'https://')):
+                    c['url'] = f"https://www.shl.com/{str(c['url']).lstrip('/')}"
+        
+        result = rerank(query, candidates)
+        
+        # Post-process result to ensure URLs are properly formatted
+        if 'recommended_assessments' in result:
+            result['recommended_assessments'] = validate_and_fix_urls(result['recommended_assessments'])
+            
+        return result
+    except Exception as e:
+        import traceback
+        error_details = traceback.format_exc()
+        print(f"Error: {str(e)}\n{error_details}")
+        return {"error": f"Error processing request: {str(e)}"}
+
+iface = gr.Interface(
+    fn=recommend,
+    inputs=gr.Textbox(label="Enter Job Description", lines=4),
+    outputs="json",
+    title="SHL Assessment Recommender",
+    description="Paste a job description to get the most relevant SHL assessments."
+)
+
+if __name__ == "__main__":
+    iface.launch()