Upload 2 files

app.py

@@ -1,28 +1,38 @@
-
 import gradio as gr
 import time
-from sentence_transformers import SentenceTransformer, util
 import os
 import json
+import torch
+from transformers import AutoTokenizer, AutoModel
+
+# --- Path Configuration ---
+# Get the absolute path of the directory containing this script
+script_dir = os.path.dirname(os.path.abspath(__file__))
+
+# Check if running in a Hugging Face Space
+is_hf_space = "SPACE_ID" in os.environ
 
-print(os.environ)
-# Determine model path based on environment
-if "SPACE_ID" in os.environ:
-    # Running on Hugging Face Spaces
-    # Assumes the model is in a repository with the same name as the space
-    model_path ='https://huggingface.co/philtoms/minilm-alice-base-rsft-v1'
-    print(f"Running on HF Spaces. Using model: {model_path}")
+if is_hf_space:
+    # In a Space, load model from the Hub and data from the repo root
+    model_path = os.environ.get("MODEL_REPO_ID", "philtoms/minilm-alice-base-rsft-v1")
+    data_path = "alice_pairs.jsonl"
+    print(f"Running on HF Spaces. Using model from Hub: {model_path}")
 else:
-    # Running locally
-    model_path = "../models/minilm-alice-base-rsft-v1/final"
-    print(f"Running locally. Using model: {model_path}")
+    # Locally, construct absolute paths based on the script's location
+    model_path = os.path.join(script_dir, "..", "models", "minilm-alice-base-rsft-v1", "final")
+    data_path = os.path.join(script_dir, "..", "data", "alice_pairs.jsonl")
+    print(f"Running locally. Using local model at: {model_path}")
 
-# Load the model
-model = SentenceTransformer(model_path)
+# --- Model and Tokenizer Loading ---
+try:
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    model = AutoModel.from_pretrained(model_path)
+except Exception as e:
+    raise gr.Error(f"Failed to load model from '{model_path}'. Error: {e}")
 
-# Load the dataset
-# Adjust the data path for local vs. HF environment
-data_path = "alice_pairs.jsonl" if "SPACE_ID" in os.environ else "../data/alice_pairs.jsonl"
+# --- Dataset Loading ---
+if not os.path.exists(data_path):
+     raise gr.Error(f"Data file not found at '{data_path}'. Please ensure the file exists.")
 
 dataset = []
 with open(data_path, "r") as f:
@@ -30,21 +40,28 @@ with open(data_path, "r") as f:
         dataset.append(json.loads(line))
 
 corpus = [item["passage"] for item in dataset]
-corpus_embeddings = model.encode(corpus, convert_to_tensor=True)
+
+# Pre-compute corpus embeddings
+with torch.no_grad():
+    encoded_corpus = tokenizer(corpus, padding=True, truncation=True, return_tensors='pt')
+    corpus_embeddings = model(**encoded_corpus).last_hidden_state.mean(dim=1)
 
 def find_similar(prompt, top_k):
     start_time = time.time()
-    
-    prompt_embedding = model.encode(prompt, convert_to_tensor=True)
-    cos_scores = util.cos_sim(prompt_embedding, corpus_embeddings)[0]
-    top_results = cos_scores.topk(k=int(top_k))
-    
+
+    with torch.no_grad():
+        encoded_prompt = tokenizer(prompt, padding=True, truncation=True, return_tensors='pt')
+        prompt_embedding = model(**encoded_prompt).last_hidden_state.mean(dim=1)
+        
+        cos_scores = torch.nn.functional.cosine_similarity(prompt_embedding, corpus_embeddings, dim=1)
+        top_results = torch.topk(cos_scores, k=int(top_k))
+
     end_time = time.time()
-    
+
     results = []
-    for score, idx in zip(top_results[0], top_results[1]):
+    for score, idx in zip(top_results.values, top_results.indices):
         results.append((corpus[idx], score.item()))
-        
+
     return results, f"{(end_time - start_time) * 1000:.2f} ms"
 
 iface = gr.Interface(
@@ -57,9 +74,9 @@ iface = gr.Interface(
         gr.Dataframe(headers=["Response", "Score"]),
         gr.Textbox(label="Time Taken")
     ],
-    title="RSFT Alice embeddings",
-    description="Enter a prompt and get the most similar sentences from the corpus."
+    title="RSFT Alice Embeddings (Transformers)",
+    description=f"Enter a prompt to find similar sentences from the corpus."
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()

requirements.txt

@@ -1,2 +1,4 @@
 gradio
-sentence-transformers
+transformers
+torch
+huggingface_hub