Update app.py

app.py

@@ -28,7 +28,7 @@ class Settings:
     AGGREGATOR_URL = os.getenv("AGGREGATOR_URL", "http://192.168.1.104:8002")
     
     # Model settings
-    MODEL_REPO = "https://huggingface.co/microsoft/florence-2-large"
+    MODEL_REPO = "https://huggingface.co/facebook/opt-125m"
     
     # Server settings
     TENSOR_SERVER_TIMEOUT = 30  # seconds
@@ -232,16 +232,30 @@ async def split_model_weights():
                     raise Exception(f"Failed to process chunk {chunk_id} at offset {start_pos}: {str(e)}")
                 
                 # Create chunk metadata
+                # Assign vocab_offset based on cumulative sizes of earlier chunks
+                # so that chunks map to disjoint vocab ranges for aggregation.
+                cumulative = 0
+                for cid, c in state.model_chunks.items():
+                    try:
+                        cumulative += int(c.config.get('shard_dim', c.config.get('size_bytes', 1)))
+                    except Exception:
+                        cumulative += 1
+
+                cfg = {
+                    "start_offset": start_pos,
+                    "size_bytes": current_chunk_size,
+                    "is_last_chunk": chunk_id == num_chunks - 1,
+                    "total_chunks": num_chunks,
+                    "original_file": os.path.basename(model_file),
+                    # minimal shard mapping; users should adjust shard_dim to real local vocab size
+                    "vocab_offset": cumulative,
+                    "shard_dim": int(current_chunk_size) if current_chunk_size > 0 else 1
+                }
+
                 state.model_chunks[chunk_id] = ModelChunk(
                     chunk_id=chunk_id,
                     files=[f"chunk_{chunk_id}.bin"],
-                    config={
-                        "start_offset": start_pos,
-                        "size_bytes": current_chunk_size,
-                        "is_last_chunk": chunk_id == num_chunks - 1,
-                        "total_chunks": num_chunks,
-                        "original_file": os.path.basename(model_file)
-                    },
+                    config=cfg,
                     size_bytes=current_chunk_size,
                     status="ready"
                 )
@@ -331,7 +345,10 @@ async def split_model_weights():
                         "size_bytes": chunk_total_size,
                         "num_parameters": sum(weights[k].nelement() for k in current_chunk),
                         "input_size": weights[current_chunk[0]].size(1) if len(current_chunk) > 0 else 0,
-                        "output_size": weights[current_chunk[-1]].size(0) if len(current_chunk) > 0 else 0
+                        "output_size": weights[current_chunk[-1]].size(0) if len(current_chunk) > 0 else 0,
+                        # assign a vocab_offset cumulatively
+                        "vocab_offset": sum(int(c.config.get('shard_dim', c.config.get('size_bytes', 1))) for c in state.model_chunks.values()),
+                        "shard_dim": int(chunk_total_size)
                     }
                 )
                 
@@ -1130,9 +1147,11 @@ async def startup_event():
                 )
                 distribution_tasks.append(task)
                 print(f"[INFO] Sending chunk {chunk_id} to {server_url}")
-                
                 # Track assignments for future reference
-                chunk.server_assignments.append(server_url)
+                try:
+                    chunk.server_assignments.append(server_url)
+                except Exception:
+                    pass
             
             if distribution_tasks:
                 print(f"[INFO] Distributing {len(distribution_tasks)} chunks...")