updated files

Dockerfile

@@ -1,25 +1,55 @@
-# Dockerfile.muril
+# Dockerfile — CPU-optimized, robust for Hugging Face Spaces / local deployment
 FROM python:3.11-slim
 
 ENV DEBIAN_FRONTEND=noninteractive
-RUN apt-get update && apt-get install -y --no-install-recommends \
-    build-essential git curl ca-certificates \
- && rm -rf /var/lib/apt/lists/*
-
 WORKDIR /app
 
+# ---------------------------
+# System deps
+# ---------------------------
+RUN apt-get update \
+ && apt-get install -y --no-install-recommends \
+    build-essential git curl ca-certificates libsndfile1 \
+ && rm -rf /var/lib/apt/lists/*
+
+# ---------------------------
+# Copy requirements files
+# ---------------------------
+# We'll install torch (CPU) first from the official PyTorch index, then install other requirements.
 COPY requirements.txt /app/requirements.txt
-RUN pip install --no-cache-dir -r /app/requirements.txt
 
-# Copy app
-COPY main.py /app/main.py
+# ---------------------------
+# Install PyTorch CPU wheel first to avoid torchvision/torch mismatch
+# (adjust CPU wheel index if you want GPU/CUDA variant)
+# ---------------------------
+RUN pip --no-cache-dir install --upgrade pip setuptools wheel \
+ && pip --no-cache-dir install --index-url https://download.pytorch.org/whl/cpu \
+      "torch>=2.1.0" "torchvision>=0.16.0" "torchaudio>=2.1.0" \
+ && pip --no-cache-dir install --no-deps -r /app/requirements.txt
 
-ENV HF_HOME=/app/hf_cache
-ENV TRANSFORMERS_CACHE=/app/hf_cache
-ENV TORCH_DISABLE_CUDA=1
-ENV OUT_DIR=/app/export_artifacts
-ENV MODEL_DIR=/app/muril_multilang_out
+# ---------------------------
+# Copy app code
+# ---------------------------
+COPY . /app
+
+# Make entrypoint executable
+RUN chmod +x /app/entrypoint.sh || true
+
+# Create export artifacts dir (mountable)
+RUN mkdir -p /app/export_artifacts
+
+# Default envs (override at runtime)
+ENV HF_REPO="Sp2503/Finetuned-multilingualdataset-MuriL-model"
+ENV MODEL_DIR="$HF_REPO"
+ENV CSV_PATH="/app/export_artifacts/muril_multilingual_dataset.csv"
+ENV OUT_EMBED_PATH="/app/export_artifacts/answer_embeddings.pt"
+ENV HF_CACHE_DIR="/app/hf_cache"
+ENV UPLOAD_BACK="false"
+ENV FORCE_REGEN="false"
 ENV PORT=7860
+ENV DEVICE="cpu"
+
+EXPOSE ${PORT}
 
-EXPOSE 7860
-CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]
+# ENTRYPOINT: run regeneration (if required) and start uvicorn
+ENTRYPOINT ["/app/entrypoint.sh"]

entrypoint.sh

@@ -0,0 +1,38 @@
+#!/usr/bin/env bash
+set -e
+
+CSV_PATH=${CSV_PATH:-/app/export_artifacts/muril_multilingual_dataset.csv}
+EMBED_PATH=${OUT_EMBED_PATH:-/app/export_artifacts/answer_embeddings.pt}
+MODEL_DIR=${MODEL_DIR:-$HF_REPO}
+HF_REPO=${HF_REPO:-Sp2503/Finetuned-multilingualdataset-MuriL-model}
+FORCE_REGEN=${FORCE_REGEN:-false}
+UPLOAD_BACK=${UPLOAD_BACK:-false}
+
+echo "Entrypoint: CSV_PATH=$CSV_PATH EMBED_PATH=$EMBED_PATH MODEL_DIR=$MODEL_DIR FORCE_REGEN=$FORCE_REGEN"
+
+mkdir -p "$(dirname "$EMBED_PATH")"
+
+need_regen=false
+if [ ! -f "$EMBED_PATH" ]; then
+  echo "Embeddings not found -> will regenerate."
+  need_regen=true
+fi
+
+if [ "$FORCE_REGEN" = "true" ] || [ "$FORCE_REGEN" = "1" ]; then
+  echo "FORCE_REGEN -> will regenerate embeddings."
+  need_regen=true
+fi
+
+if [ "$need_regen" = "true" ]; then
+  export MODEL_DIR="${MODEL_DIR}"
+  export CSV_PATH="${CSV_PATH}"
+  export OUT_EMBED_PATH="${EMBED_PATH}"
+  export HF_REPO="${HF_REPO}"
+  export UPLOAD_BACK="${UPLOAD_BACK}"
+  python /app/regenerate_embeddings.py
+  echo "Regeneration finished."
+else
+  echo "Skipping regeneration."
+fi
+
+exec uvicorn main:app --host 0.0.0.0 --port ${PORT:-7860}

model_loader.py

@@ -0,0 +1,78 @@
+# model_loader.py
+import os
+from pathlib import Path
+import torch
+
+def load_model_and_tokenizer(model_repo_dir_or_local_path: str, base_model_id: str = "google/muril-base-cased", device: str = None):
+    """
+    Robust loader that:
+      - loads full model if model.safetensors or pytorch_model.bin exists in model_repo_dir_or_local_path
+      - otherwise loads base_model_id then applies PEFT adapter from adapter_model.safetensors (if present)
+    Returns (tokenizer, model_on_device, backend_str)
+    """
+    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+    model_dir = Path(model_repo_dir_or_local_path)
+
+    # prefer safetensors full model if present
+    full_model_files = ["model.safetensors", "pytorch_model.bin"]
+    adapter_files = ["adapter_model.safetensors", "adapter_config.json", "adapter.safetensors"]
+
+    # Import delayed to allow environment to control package errors
+    from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification
+
+    tokenizer = None
+    model = None
+
+    # 1) Try to load full model from model_dir
+    for f in full_model_files:
+        full_path = model_dir / f
+        if full_path.exists():
+            tokenizer = AutoTokenizer.from_pretrained(str(model_dir), use_fast=True)
+            try:
+                model = AutoModel.from_pretrained(str(model_dir), trust_remote_code=False)
+                backend = "full-AutoModel"
+            except Exception:
+                model = AutoModelForSequenceClassification.from_pretrained(str(model_dir), trust_remote_code=False)
+                backend = "full-AutoModelForSequenceClassification"
+            model.to(device)
+            model.eval()
+            return tokenizer, model, backend
+
+    # 2) If full model not present, check for adapter files
+    has_adapter = any((model_dir / af).exists() for af in adapter_files)
+    if has_adapter:
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(str(model_dir), use_fast=True)
+        except Exception:
+            tokenizer = AutoTokenizer.from_pretrained(base_model_id, use_fast=True)
+
+        try:
+            base = AutoModel.from_pretrained(str(model_dir))
+            base_name = str(model_dir)
+        except Exception:
+            base = AutoModel.from_pretrained(base_model_id)
+            base_name = base_model_id
+
+        base.to(device)
+        base.eval()
+
+        try:
+            from peft import PeftModel
+            adapter_dir = str(model_dir)
+            # If adapter files exist, PeftModel.from_pretrained will pick them up
+            peft_model = PeftModel.from_pretrained(base, adapter_dir, is_trainable=False)
+            peft_model.to(device)
+            peft_model.eval()
+            return tokenizer, peft_model, "peft-attached"
+        except Exception as e:
+            raise RuntimeError(f"Failed to load/apply PEFT adapter from {model_dir}: {e}") from e
+
+    # 3) Fallback: try direct load (may fail)
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(str(model_dir), use_fast=True)
+        model = AutoModel.from_pretrained(str(model_dir))
+        model.to(device)
+        model.eval()
+        return tokenizer, model, "auto-fallback"
+    except Exception as e:
+        raise RuntimeError(f"Unable to load model or adapters from {model_dir}. Error: {e}") from e

regenerate_embeddings.py

@@ -0,0 +1,125 @@
+# regenerate_embeddings.py
+"""
+Regenerate answer embeddings using the MuRIL model.
+This script:
+ - downloads model (if MODEL_DIR is a repo id),
+ - reads CSV at CSV_PATH,
+ - computes mean-pooled, L2-normalized embeddings for 'answer' column,
+ - saves embeddings to OUT_EMBED_PATH.
+
+Exit codes:
+ - 0 on success
+ - non-zero on failure
+"""
+import os, argparse, math, sys
+from pathlib import Path
+import torch
+import pandas as pd
+from tqdm.auto import tqdm
+from transformers import AutoTokenizer, AutoModel
+from huggingface_hub import snapshot_download
+
+def mean_pooling(last_hidden_state, attention_mask):
+    input_mask_expanded = attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float()
+    sum_embeddings = torch.sum(last_hidden_state * input_mask_expanded, 1)
+    sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
+    return sum_embeddings / sum_mask
+
+def parse_env():
+    # ENV-friendly arg parsing
+    cfg = {}
+    cfg['model_dir'] = os.getenv("MODEL_DIR", os.getenv("HF_REPO", "Sp2503/Finetuned-multilingualdataset-MuriL-model"))
+    cfg['csv_path'] = os.getenv("CSV_PATH", "/app/export_artifacts/muril_multilingual_dataset.csv")
+    cfg['out_path'] = os.getenv("OUT_EMBED_PATH", "/app/export_artifacts/answer_embeddings.pt")
+    cfg['batch_size'] = int(os.getenv("EMBED_BATCH_SIZE", "64"))
+    cfg['device'] = os.getenv("DEVICE", "cuda" if torch.cuda.is_available() else "cpu")
+    cfg['download_cache'] = os.getenv("HF_CACHE_DIR", "/tmp/hf_cache")
+    cfg['upload_back'] = os.getenv("UPLOAD_BACK", "false").lower() in ("1","true","yes")
+    cfg['hf_repo'] = os.getenv("HF_REPO", None)  # used for upload_back if set
+    return cfg
+
+def main():
+    cfg = parse_env()
+    print("Regenerate embeddings with config:", cfg)
+    model_dir = cfg['model_dir']
+    # If model_dir looks like a HF repo id (contains '/'), snapshot_download to local cache
+    if "/" in model_dir and not os.path.isdir(model_dir):
+        print("Detected HF repo id for model. snapshot_download ->", cfg['download_cache'])
+        try:
+            model_dir = snapshot_download(repo_id=cfg['model_dir'], repo_type="model", cache_dir=cfg['download_cache'])
+            print("Downloaded model to:", model_dir)
+        except Exception as e:
+            print("Failed to snapshot_download model:", e, file=sys.stderr)
+            sys.exit(2)
+
+    csv_path = cfg['csv_path']
+    out_path = cfg['out_path']
+    batch_size = cfg['batch_size']
+    device = cfg['device']
+    print(f"Loading CSV: {csv_path}")
+    if not os.path.isfile(csv_path):
+        print(f"CSV not found at {csv_path}", file=sys.stderr)
+        sys.exit(3)
+    df = pd.read_csv(csv_path, dtype=str).fillna("")
+    if 'answer' not in df.columns:
+        print("CSV must contain 'answer' column", file=sys.stderr)
+        sys.exit(4)
+    answers = df['answer'].astype(str).tolist()
+    print(f"Encoding {len(answers)} answers on device {device} (batch_size={batch_size})")
+
+    # Load tokenizer & model
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_dir, use_fast=True)
+        model = AutoModel.from_pretrained(model_dir)
+        model.to(device)
+        model.eval()
+    except Exception as e:
+        print("Failed to load model/tokenizer:", e, file=sys.stderr)
+        sys.exit(5)
+
+    # compute embeddings
+    all_embs = []
+    try:
+        with torch.inference_mode():
+            for i in tqdm(range(0, len(answers), batch_size), desc="Batches"):
+                batch = answers[i:i+batch_size]
+                enc = tokenizer(batch, padding=True, truncation=True, max_length=256, return_tensors="pt")
+                input_ids = enc["input_ids"].to(device)
+                attention_mask = enc["attention_mask"].to(device)
+                out = model(input_ids=input_ids, attention_mask=attention_mask, return_dict=True)
+                pooled = mean_pooling(out.last_hidden_state, attention_mask)      # (B, H)
+                pooled = torch.nn.functional.normalize(pooled, p=2, dim=1)        # L2-normalize
+                all_embs.append(pooled.cpu())
+    except Exception as e:
+        print("Error during encoding:", e, file=sys.stderr)
+        sys.exit(6)
+
+    all_embs = torch.cat(all_embs, dim=0)
+    print("Final embeddings shape:", all_embs.shape)
+    Path(out_path).parent.mkdir(parents=True, exist_ok=True)
+    torch.save(all_embs, out_path)
+    print("Saved embeddings to:", out_path)
+
+    # Optional: upload back to HF repo (requires HF_TOKEN set and HF_REPO)
+    if cfg['upload_back'] and cfg['hf_repo']:
+        try:
+            from huggingface_hub import HfApi
+            api = HfApi()
+            print(f"Uploading {out_path} back to repo {cfg['hf_repo']} ...")
+            api.upload_file(
+                path_or_fileobj=out_path,
+                path_in_repo=os.path.basename(out_path),
+                repo_id=cfg['hf_repo'],
+                repo_type="model",
+            )
+            print("Upload complete.")
+        except Exception as e:
+            print("Upload back failed:", e, file=sys.stderr)
+
+    # quick sanity check
+    norms = (all_embs * all_embs).sum(dim=1)
+    print("Sample norms (should be ~1.0):", norms[:5].tolist())
+    return 0
+
+if __name__ == "__main__":
+    sys.exit(main())

requirements.txt

@@ -5,3 +5,5 @@ torch>=2.1.0
 transformers==4.46.0
 huggingface_hub>=0.14.1
 tqdm
+peft
+safetensors