major update: all data saved on HF (prompts, results), unified utilities

src/config/settings.py

@@ -17,9 +17,9 @@ DEFAULT_SIMILARITY_MODEL = "all-MiniLM-L6-v2"
 # Token limits
 MAX_USER_TOKENS = 5
 
-# Data file paths
-PROMPTS_DATA_PATH = "data/prompts.jsonl"
-RESULTS_DATA_PATH = "data/results.jsonl"
+# Huggingface data repositories
+HF_PROMPTS_REPO = "alon-albalak/collaborative-decoding-prompts"
+HF_RESULTS_REPO = "alon-albalak/collaborative-decoding-results"
 
 # Server configuration
 DEFAULT_SERVER_NAME = "127.0.0.1"

src/models/data_manager.py

@@ -1,22 +1,21 @@
 """Data loading and saving functionality"""
 
 import json
-import os
 import random
 import datetime
 import uuid
 from pathlib import Path
-from datasets import load_dataset
 from huggingface_hub import CommitScheduler
 
-HF_REPO_ID = "alon-albalak/collaborative-decoding-results"
+from src.config.settings import HF_RESULTS_REPO, HF_PROMPTS_REPO
+from src.utils.hf_data_manager import HFDataManager
 
-JSON_DATASET_DIR = Path("results")
+JSON_DATASET_DIR = Path("testing/data/results")
 JSON_DATASET_DIR.mkdir(parents=True, exist_ok=True)
 JSON_DATASET_PATH = JSON_DATASET_DIR / f"results_{uuid.uuid4()}.json"
 
 scheduler = CommitScheduler(
-    repo_id=HF_REPO_ID,
+    repo_id=HF_RESULTS_REPO,
     repo_type="dataset",
     folder_path=JSON_DATASET_DIR.as_posix(),
     path_in_repo="data",
@@ -30,10 +29,11 @@ class DataManager:
         self.prompts_data = []
         self.results = None
 
-    def load_prompts_data(self, filepath="data/prompts.jsonl"):
-        """Load prompts data from JSONL file"""
-        with open(filepath, "r") as f:
-            self.prompts_data = [json.loads(line) for line in f]
+    def load_prompts_data(self):
+        """Load prompts data"""
+        self.prompts_data = self.load_from_hf(HF_PROMPTS_REPO)
+        if not self.prompts_data:
+            raise RuntimeError("No prompts data loaded from Hugging Face.")
     
     def get_random_prompt(self):
         """Get a random prompt from loaded data"""
@@ -44,7 +44,7 @@ class DataManager:
     def get_results(self):
         """Get all results data, loading if not already loaded."""
         if self.results is None:
-            self.results = self.load_results_from_hf()
+            self.results = self.load_from_hf(HF_RESULTS_REPO)
         return self.results
     
     def add_results(self, new_results):
@@ -53,14 +53,9 @@ class DataManager:
             raise RuntimeError("Results not loaded. Call get_results() first.")
         self.results.extend(new_results)
     
-    def load_results_from_hf(self, hf_repo=HF_REPO_ID):
-        """Load results data from Hugging Face dataset repository."""
-        try:
-            dataset = load_dataset(hf_repo, split="train")
-            return dataset.to_list()
-        except Exception as e:
-            print(f"Error loading dataset from Hugging Face: {e}")
-            return []
+    def load_from_hf(self, hf_repo):
+        """Load data from Hugging Face dataset repository."""
+        return HFDataManager.load_from_hf(hf_repo)
     
     def save_interaction_to_hf(self, prompt_data, user_continuation, generated_response, 
                                cosine_distance, session_id, num_user_tokens):

src/models/llm_manager.py

@@ -2,7 +2,8 @@
 
 import os
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from src.utils.model_loader import ModelLoader
+from src.utils.prompt_formatter import PromptFormatter
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
@@ -12,28 +13,12 @@ class LLMManager:
     def __init__(self):
         self.model = None
         self.tokenizer = None
-
-        if torch.cuda.is_available():
-            device = "cuda"
-            dtype = torch.float16
-        elif torch.backends.mps.is_available():
-            device = "mps"
-            dtype = torch.float16
-        else:
-            device = "cpu"
-            dtype = torch.float32
-
-        self.device = device
-        self.dtype = dtype
+        self.device, self.dtype = ModelLoader.get_device_and_dtype()
         
     def load_models(self, model_name="meta-llama/Llama-3.2-1B-Instruct"):
         """Load the LLM model and tokenizer"""
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.model = AutoModelForCausalLM.from_pretrained(model_name, dtype=self.dtype, low_cpu_mem_usage=True)
-        self.model = self.model.to(self.device)
-        
-        if self.tokenizer.pad_token is None:
-            self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.model_name = model_name
+        self.model, self.tokenizer, self.device, self.dtype = ModelLoader.load_model_and_tokenizer(model_name)
     
     def validate_user_input(self, user_input, max_tokens=5):
         """Validate that user input is within token limits"""
@@ -63,13 +48,22 @@ class LLMManager:
             token_texts.append(token_text)
         return tokens, token_texts
     
+    def extract_assistant_response(self, full_response: str) -> str:
+        """Extract the assistant's response from the full generated text"""
+        return PromptFormatter.extract_assistant_response(self.model_name, full_response)
+
+    def format_prompt(self, prompt: str, partial_response: str, continuation: str) -> str:
+        """Format the full prompt for generation"""
+        return PromptFormatter.format_prompt(self.model_name, prompt, partial_response, continuation)
+
+
     def generate_response_from_user_input(self, prompt, partial_response, user_continuation):
         """Generate a full response from user's continuation"""
         if not self.model or not self.tokenizer:
             raise RuntimeError("Models not loaded. Call load_models() first.")
         
         # TODO: make this more robust for multiple models, needs to be formatted correctly
-        full_prompt = f"{prompt}\n\nAssistant: {partial_response}{user_continuation}"
+        full_prompt = self.format_prompt(prompt, partial_response, user_continuation)
         
         inputs = self.tokenizer(full_prompt, return_tensors="pt", padding=True, truncation=True)
         inputs = {k: v.to(self.device) for k, v in inputs.items()}
@@ -86,5 +80,6 @@ class LLMManager:
             )
         
         full_response = self.tokenizer.decode(outputs[0].cpu(), skip_special_tokens=True)
-        assistant_part = full_response.split("Assistant: ")[-1]
+        
+        assistant_part = self.extract_assistant_response(full_response)
         return assistant_part

src/scoring/scorer.py

@@ -33,12 +33,15 @@ class Scorer:
         else:
             return "Hard", user_token_count
     
-    def calculate_rank_and_percentile(self, user_score, prompt_results, user_tokens):
+    def calculate_rank_and_percentile(self, user_score, prompt_results, user_tokens, separate_by_token_count=False):
         """Calculate user's rank and percentile among users with same prompt and token count."""
         # Filter to only same prompt and same token count
-        comparable_scores = [r["cosine_distance"] for r in prompt_results 
-                           if r["num_user_tokens"] == user_tokens]
-        
+        if separate_by_token_count:
+            comparable_scores = [r["cosine_distance"] for r in prompt_results 
+                            if r["num_user_tokens"] == user_tokens]
+        else:
+            comparable_scores = [r["cosine_distance"] for r in prompt_results]
+
         if not comparable_scores:
             return None, None
         
@@ -117,7 +120,7 @@ class Scorer:
             html_content += f"""
                 <div class="score-metric">
                     <div class="metric-value">#{rank}</div>
-                    <div class="metric-label">Rank out of {same_category_attempts+1}</div>
+                    <div class="metric-label">Rank out of {same_category_attempts}</div>
                 </div>
                 <div class="score-metric">
                     <div class="metric-value">{percentile:.1f}%</div>

src/ui/page_handlers.py

@@ -404,8 +404,7 @@ class PageHandlers:
         
         # Create enhanced score display with progress bars and metrics
         user_tokens = self.app.llm_manager.count_tokens(user_text)
-        same_category_attempts = len([r for r in prompt_results 
-                                    if r["num_user_tokens"] == user_tokens]) if prompt_results else 1
+        same_category_attempts = len(prompt_results)
         
         score_text = self.app.scorer.create_enhanced_score_display(
             cosine_distance, rank, percentile, user_tokens, same_category_attempts

src/utils/hf_data_manager.py

@@ -0,0 +1,43 @@
+"""Unified HuggingFace dataset operations"""
+
+import json
+from pathlib import Path
+from datasets import Dataset, load_dataset
+
+class HFDataManager:
+    """Handles all HuggingFace dataset loading and saving operations"""
+    
+    @staticmethod
+    def load_from_hf(hf_repo):
+        """Load data from HuggingFace dataset repository"""
+        try:
+            dataset = load_dataset(hf_repo, split="train")
+            return dataset.to_list()
+        except Exception:
+            # Return empty list if dataset doesn't exist or can't be loaded
+            return []
+    
+    @staticmethod
+    def push_to_hf(data, repo_id, private=True):
+        """Push data to HuggingFace dataset repository"""
+        dataset = Dataset.from_list(data)
+        dataset.push_to_hub(repo_id, private=private)
+    
+    @staticmethod
+    def save_to_jsonl(data, file_path):
+        """Save data to local JSONL file"""
+        file_path = Path(file_path)
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+        
+        with open(file_path, "w") as f:
+            for item in data:
+                f.write(json.dumps(item) + "\n")
+    
+    @staticmethod
+    def load_from_jsonl(file_path):
+        """Load data from local JSONL file"""
+        data = []
+        with open(file_path, "r") as f:
+            for line in f:
+                data.append(json.loads(line.strip()))
+        return data

src/utils/model_loader.py

@@ -0,0 +1,39 @@
+"""Unified model loading and device management"""
+
+import os
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+class ModelLoader:
+    """Handles device detection and model/tokenizer loading"""
+    
+    @staticmethod
+    def get_device_and_dtype():
+        """Determine the best available device and dtype"""
+        if torch.cuda.is_available():
+            return "cuda", torch.float16
+        elif torch.backends.mps.is_available():
+            return "mps", torch.float16
+        else:
+            return "cpu", torch.float32
+    
+    @staticmethod
+    def load_model_and_tokenizer(model_name="meta-llama/Llama-3.2-1B-Instruct"):
+        """Load model and tokenizer with optimal device/dtype settings"""
+        device, dtype = ModelLoader.get_device_and_dtype()
+        
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name, 
+            dtype=dtype, 
+            low_cpu_mem_usage=True
+        )
+        model = model.to(device)
+        
+        # Set pad token if needed
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+        
+        return model, tokenizer, device, dtype

src/utils/prompt_formatter.py

@@ -0,0 +1,28 @@
+"""Unified prompt formatting and response extraction"""
+
+import re
+
+class PromptFormatter:
+    """Handles prompt formatting and assistant response extraction"""
+    
+    @staticmethod
+    def format_prompt(model_name, prompt, partial_response, continuation):
+        """Format the full prompt for generation"""
+        if "meta-llama" in model_name:
+            return f"{prompt}\n\nAssistant: {partial_response}{continuation}"
+        else:
+            raise NotImplementedError(f"Prompt formatting not implemented for model: {model_name}")
+    
+    @staticmethod
+    def extract_assistant_response(model_name, full_response):
+        """Extract the assistant's response from the full generated text"""
+        if "meta-llama" in model_name:
+            # Check if we have multiple assistant tags and get the last one
+            assistant_tags = re.findall(r"Assistant:\s*", full_response)
+            if len(assistant_tags) > 1:
+                print(f"Found multiple assistant tags ({len(assistant_tags)})\nFull response:\n{full_response}\n**")
+            
+            # Only split on the first assistant tag
+            return full_response.split("Assistant:", maxsplit=1)[-1].strip()
+        else:
+            raise NotImplementedError(f"Response extraction not implemented for model: {model_name}")