siddhm11 commited on
Commit
fec1e1a
·
verified ·
1 Parent(s): 49c43e6

add: quick loading snippet for Python users

Browse files
Files changed (1) hide show
  1. load_model.py +59 -0
load_model.py ADDED
@@ -0,0 +1,59 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Quick-start: Load and use the ResearchIT Phase 6 reranker.
3
+
4
+ Usage:
5
+ python load_model.py
6
+
7
+ Or import in your code:
8
+ from load_model import load_reranker, predict_scores
9
+ """
10
+ import numpy as np
11
+
12
+ def load_reranker(model_path: str = "production_model/reranker_v1.txt"):
13
+ """Load the LightGBM reranker model."""
14
+ import lightgbm as lgb
15
+ model = lgb.Booster(model_file=model_path)
16
+ assert model.num_feature() == 37, f"Expected 37 features, got {model.num_feature()}"
17
+ return model
18
+
19
+ def predict_scores(model, features: np.ndarray) -> np.ndarray:
20
+ """
21
+ Predict reranking scores for candidates.
22
+
23
+ Args:
24
+ model: LightGBM Booster
25
+ features: (N, 37) float32 array — see feature_schema.json for column order
26
+
27
+ Returns:
28
+ (N,) float64 array — higher score = more relevant
29
+ """
30
+ assert features.shape[1] == 37, f"Expected 37 features, got {features.shape[1]}"
31
+ return model.predict(features)
32
+
33
+ # Feature schema (must match this exact order)
34
+ FEATURE_SCHEMA = [
35
+ "qdrant_cosine_score", "candidate_position", "candidate_citation_count",
36
+ "candidate_log_citations", "candidate_influential_citations",
37
+ "candidate_age_days", "candidate_recency_score", "query_citation_count",
38
+ "query_age_days", "year_diff", "same_primary_category", "co_citation_count",
39
+ "shared_author_count", "candidate_is_newer", "query_log_citations",
40
+ "citation_count_ratio", "age_ratio", "candidate_citations_per_year",
41
+ "query_num_references", "candidate_num_cited_by",
42
+ "ewma_longterm_similarity", "ewma_shortterm_similarity",
43
+ "ewma_negative_similarity", "cluster_importance",
44
+ "cluster_distance_to_medoid", "is_suppressed_category",
45
+ "onboarding_category_match", "user_total_saves", "user_total_dismissals",
46
+ "user_days_since_last_save", "user_session_save_count",
47
+ "cosine_x_recency", "cosine_x_citations", "category_x_recency",
48
+ "cosine_x_cocitation", "position_inverse", "citations_x_recency",
49
+ ]
50
+
51
+ if __name__ == "__main__":
52
+ model = load_reranker()
53
+ print(f"Model loaded: {model.num_trees()} trees, {model.num_feature()} features")
54
+
55
+ # Test with dummy input
56
+ dummy = np.zeros((10, 37), dtype=np.float32)
57
+ scores = predict_scores(model, dummy)
58
+ print(f"Dummy scores: {scores[:5]}")
59
+ print("✅ Model works!")