Update app.py

app.py

@@ -373,14 +373,28 @@ def train_model(
     model_repo_name, 
     dataset_repo_name, 
     epochs=1, 
-    batch_size=4,  # Increased for A100
-    grad_accum_steps=4,
+    batch_size=4,
+    grad_accum_steps=2,
     learning_rate=2e-4,
     progress=gr.Progress()
 ):
     progress(0, desc="Setting up environment...")
     log = []
     
+    # Install sentencepiece if it's not already installed
+    progress(0.02, desc="Installing required dependencies...")
+    try:
+        import sentencepiece
+        log.append("SentencePiece already installed")
+    except ImportError:
+        log.append("Installing SentencePiece...")
+        try:
+            subprocess.check_call([sys.executable, "-m", "pip", "install", "sentencepiece"])
+            log.append("SentencePiece installed successfully")
+        except Exception as e:
+            log.append(f"Error installing SentencePiece: {e}")
+            # Continue anyway, we'll try other tokenizer approaches if this fails
+    
     # Clean up any existing model files to save space
     if os.path.exists("./model_files"):
         try:
@@ -508,45 +522,66 @@ def train_model(
             torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
         )
         
-        # --- Load Tokenizer (from a compatible model) ---
-        # Following the pattern from reference code
+        # --- Load Tokenizer (using public models) ---
         progress(0.3, desc="Loading tokenizer...")
         
         # Try to load a compatible tokenizer
         try:
-            # First try loading from standard Llama 3 model
+            # First try TinyLlama which is open and uses Llama tokenizer
             tokenizer = AutoTokenizer.from_pretrained(
-                "meta-llama/Llama-3-8B",  # Using standard Llama 3 tokenizer
+                "TinyLlama/TinyLlama-1.1B-Chat-v1.0",  # Public model with Llama tokenizer
                 padding_side="right",
                 use_fast=True,
-                trust_remote_code=True
             )
-            log.append("Loaded tokenizer from meta-llama/Llama-3-8B")
+            log.append("Loaded tokenizer from TinyLlama/TinyLlama-1.1B-Chat-v1.0")
         except Exception as e1:
-            log.append(f"Couldn't load Llama-3 tokenizer: {e1}")
+            log.append(f"Couldn't load TinyLlama tokenizer: {e1}")
             try:
-                # Fallback to Llama 2
+                # Try Phi-2 which is also public
                 tokenizer = AutoTokenizer.from_pretrained(
-                    "meta-llama/Llama-2-7b-hf",
+                    "microsoft/phi-2",
                     padding_side="right",
                     use_fast=True
                 )
-                log.append("Loaded Llama-2 tokenizer as fallback")
+                log.append("Loaded Phi-2 tokenizer as fallback")
             except Exception as e2:
-                log.append(f"Couldn't load Llama-2 tokenizer: {e2}")
-                # Final fallback
-                from transformers import LlamaTokenizer
-                tokenizer = LlamaTokenizer.from_pretrained(
-                    "hf-internal-testing/llama-tokenizer",
-                    padding_side="right"
-                )
-                log.append("Loaded testing Llama tokenizer as final fallback")
+                log.append(f"Couldn't load Phi-2 tokenizer: {e2}")
+                try:
+                    # Try CodeLlama which is popular and public
+                    tokenizer = AutoTokenizer.from_pretrained(
+                        "codellama/CodeLlama-7b-hf",
+                        padding_side="right"
+                    )
+                    log.append("Loaded CodeLlama tokenizer as fallback")
+                except Exception as e3:
+                    log.append(f"Couldn't load any standard tokenizers. Using a basic tokenizer instead.")
+                    
+                    # Create a minimal tokenizer that works with our format
+                    # Assuming the vocab size is 2048 (from the RVQ token count)
+                    from transformers import PreTrainedTokenizerFast
+                    
+                    # Create a very basic tokenizer
+                    tokenizer = PreTrainedTokenizerFast(
+                        tokenizer_file=None,  # No file needed
+                        bos_token="<s>",
+                        eos_token="</s>",
+                        unk_token="<unk>",
+                        pad_token="<pad>",
+                        model_max_length=2048  # Safe default value
+                    )
+                    
+                    # Add vocabulary - creating a minimal vocab for the RVQ tokens
+                    vocab = {f"<token_{i}>": i for i in range(model.config.vocab_size)}
+                    tokenizer.add_tokens(list(vocab.keys()))
+                    log.append(f"Created basic tokenizer with {len(tokenizer)} tokens")
         
-        # Set pad token and ensure it's usable
+        # Set pad token if not already set
         if tokenizer.pad_token is None:
-            tokenizer.pad_token = tokenizer.eos_token
+            tokenizer.pad_token = tokenizer.eos_token if tokenizer.eos_token is not None else "<pad>"
+            log.append("Set pad_token to eos_token or <pad>")
         
-        log.append(f"Loaded model vocab size: {model.config.vocab_size}")
+        log.append(f"Tokenizer loaded with vocab size: {len(tokenizer)}")
+        log.append(f"Model vocab size: {model.config.vocab_size}")
         log.append(f"Input embedding shape: {model.get_input_embeddings().weight.shape}")
         
         # --- QLoRA Preparation ---