Update utils/xlnet_model.py

utils/xlnet_model.py

@@ -1,48 +1,22 @@
-import requests, tqdm, torch, numpy as np
-from torch import nn
 import os
-os.environ["HF_HOME"] = "/tmp/huggingface"
+import torch
+import numpy as np
+from torch import nn
 from transformers import XLNetModel, XLNetTokenizer
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
-
-# ------------------------------------------------------------------
 from huggingface_hub import hf_hub_download
-import torch
 
-# Downloads the file automatically from your model repo
+# Set Hugging Face cache directory
+os.environ["HF_HOME"] = "/tmp/huggingface"
+
+# Download model weights from Hugging Face Hub
 MODEL_PATH = hf_hub_download(
     repo_id="yeswanthvarma/xlnet-evaluator-model",
     filename="xlnet_answer_assessment_model.pt"
 )
 
-# Load the model
-# Load tokenizer and model architecture
-tokenizer = XLNetTokenizer.from_pretrained("xlnet-base-cased")
-model = XLNetModel.from_pretrained("xlnet-base-cased")
-
-# Then load your custom weights
-model.load_state_dict(torch.load(MODEL_PATH, map_location=torch.device("cpu")))
-model.eval()
-
-def download_model_if_needed():
-    if os.path.exists(MODEL_PATH):
-        return
-    print("▶️  Downloading XLNet weights from Hugging Face …")
-    with requests.get(HF_URL, stream=True) as r:
-        r.raise_for_status()
-        total = int(r.headers.get("content-length", 0))
-        with open(MODEL_PATH, "wb") as f, tqdm.tqdm(total=total, unit="B", unit_scale=True) as bar:
-            for chunk in r.iter_content(chunk_size=8192):
-                f.write(chunk)
-                bar.update(len(chunk))
-    print("✅  Download complete.")
-
-download_model_if_needed()
-# ------------------------------------------------------------------
-
-xlnet_available = False   # will flip to True if load succeeds
-
+# Define your custom model
 class XLNetAnswerAssessmentModel(nn.Module):
     def __init__(self):
         super().__init__()
@@ -58,6 +32,8 @@ class XLNetAnswerAssessmentModel(nn.Module):
         x = torch.relu(self.fc2(x))
         return torch.sigmoid(self.out(x))
 
+# Initialize model and tokenizer
+xlnet_available = False
 try:
     tokenizer = XLNetTokenizer.from_pretrained("xlnet-base-cased")
     model = XLNetAnswerAssessmentModel()
@@ -68,11 +44,9 @@ try:
 except Exception as e:
     print("⚠️  Could not load XLNet model → fallback to TF‑IDF\n", e)
 
-# ------------------------------------------------------------------
-# scoring helpers (unchanged)
-# ------------------------------------------------------------------
-embedding_cache = {}
-
+# -------------------------------
+# Scoring logic
+# -------------------------------
 def get_model_prediction(q, s, r):
     if not xlnet_available:
         raise ValueError("XLNet unavailable")
@@ -95,4 +69,4 @@ def get_similarity_score(q, s, r):
     try:
         return get_model_prediction(q, s, r) if xlnet_available else tfidf_similarity(s, r)
     except Exception:
-        return fallback_similarity(s, r)
+        return fallback_similarity(s, r)