Update README.md

README.md

@@ -35,24 +35,21 @@ widget:
 
 Cephalo is a series of multimodal materials science focused vision large language models (V-LLMs) designed to integrate visual and linguistic data for advanced understanding and interaction in human-AI or multi-agent AI frameworks. 
 
-A novel aspect of Cephalo's development is the innovative dataset generation method. The extraction process employs advanced algorithms to accurately detect and separate images and their corresponding textual descriptions from complex PDF documents. It involves extracting images and captions from PDFs to create well-reasoned image-text pairs, utilizing large language models (LLMs) for natural language processing. These image-text pairs are then refined and validated through LLM-based NLP processing, ensuring high-quality and contextually relevant data for training. 
-
-Cephalo can interpret complex visual scenes and generating contextually accurate language descriptions and answer queries. 
-
 The model is developed to process diverse inputs, including images and text, facilitating a broad range of applications such as image captioning, visual question answering, and multimodal content generation. The architecture combines a vision encoder model and an autoregressive transformer to process complex natural language understanding. 
 
 ![image/png](https://cdn-uploads.huggingface.co/production/uploads/623ce1c6b66fedf374859fe7/kl5GWBP9WS0D4uwd1t3S7.png)
 
 Cephalo provides a robust framework for multimodal interaction and understanding, including the development of complex generative pipelines to create 2D and 3D renderings of material microstructures as input for additive manufacturing methods.
 
-This version of Cephalo, lamm-mit/Cephalo-Idefics2-3x8b-beta, is a Mixture-of-Expert model based on the Idefics-2 model. The basic model architecture is as follows:
+This version of Cephalo, lamm-mit/Cephalo-Idefics2-3x8b-beta, is a Mixture-of-Expert model based on variants and fine-tuned versions of the Idefics-2 model. The basic model architecture is as follows:
 
 ![image/png](https://cdn-uploads.huggingface.co/production/uploads/623ce1c6b66fedf374859fe7/b7BK8ZtDzTMsyFDi0wP3w.png)
 
+The model has 20b parameters (3 experts, each 8b each, 8b active parameters during inference).
 
 ### Download Idefics-2 MoE Model and Sample inference code
 
-```markdown
+```python
 pip install transformers -U
 ```
 
@@ -69,9 +66,7 @@ def count_parameters(model):
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 model_name_moe = f"lamm-mit/Cephalo-Idefics2-3x8b-beta"
-
 config = AutoConfig.from_pretrained(model_name_moe, trust_remote_code=True)
-
 processor = AutoProcessor.from_pretrained(model_name_moe, trust_remote_code=True) 
 moe_model = AutoModelForCausalLM.from_pretrained(
     model_name_moe,config=config,
@@ -117,7 +112,7 @@ print(generated_texts)
 
 Download .py files that implement the Phi-3-V and the Mixture-of-Expert Vision model 
 
-```markdown
+```python
 pip install huggingface_hub
 ```
 
@@ -173,7 +168,7 @@ model_1 = Idefics2ForConditionalGeneration.from_pretrained(  model_id_1,
                                                            _attn_implementation="flash_attention_2", #make sure Flash Attention 2 is installed
                                                            trust_remote_code=True,
                                                            #quantization_config=quantization_config,
-                                                        )#.to (DEVICE)
+                                                        ) 
 processor = AutoProcessor.from_pretrained(
     f"{model_id_1}",
     do_image_splitting=True
@@ -186,7 +181,7 @@ processor.chat_template = IDEFICS2_CHAT_TEMPLATE
 ```
 
 Now, load the rest of the models:
-```
+```python
 model_id_2='HuggingFaceM4/idefics2-8b-chatty'
 
 model_2 = Idefics2ForConditionalGeneration.from_pretrained(  model_id_2,
@@ -194,7 +189,7 @@ model_2 = Idefics2ForConditionalGeneration.from_pretrained(  model_id_2,
                                                            _attn_implementation="flash_attention_2", #make sure Flash Attention 2 is installed
                                                            trust_remote_code=True,
                                                            #quantization_config=quantization_config,
-                                                        )#.to (DEVICE)
+                                                        ) 
 
 model_id_3='HuggingFaceM4/idefics2-8b'
 
@@ -203,38 +198,41 @@ model_3 = Idefics2ForConditionalGeneration.from_pretrained(  model_id_3,
                                                            _attn_implementation="flash_attention_2", #make sure Flash Attention 2 is installed
                                                            trust_remote_code=True,
                                                            #quantization_config=quantization_config,
-                                                        )#.to (DEVICE)
+                                                        ) 
 ```
 Put on device:
-```
+```python
 model_1.to(DEVICE)
 model_2.to(DEVICE)
 model_3.to(DEVICE)
 ```
 
 ### Construct MoE 
-```
+
+Here we show how a MoE is constructed from the set of expert models loaded earlier. We consider three models, model_1, model_2 and model_3. 
+
+```python
 dtype = torch.bfloat16  # Desired dtype for new layers
 base_model = copy.deepcopy(model_1)  # Your base model
-expert_models = [ model_1,  model_2,  model_3  ]  # List of expert models
+expert_models = [ model_1,  model_2,  model_3 ]  # List of expert models
 
 moe_config = Idefics2ForCausalLMMoEConfig(config=config, k=1, num_expert_models=len (expert_models))
-moe_model = Idefics2ForCausalLMMoE(moe_config, base_model, expert_models,  layer_dtype = dtype)#.to(device)
+moe_model = Idefics2ForCausalLMMoE(moe_config, base_model, expert_models,  layer_dtype = dtype) 
 
 count_parameters(expert_models[0]),count_parameters(moe_model)
 ```
 Delete models no longer needed:
-```
+```python
 del model_1
 del model_2
 del model_3 
 ```
 Put MoE model on device:
-```
+```python
 moe_model.to(DEVICE)
 ```
-Test if it works (untrained):
-```
+Test if it works (untrained, may not produce desirable putput since gating layers have not been trained):
+```python
 from transformers.image_utils import load_image
 
 image = load_image("https://d2r55xnwy6nx47.cloudfront.net/uploads/2018/02/Ants_Lede1300.jpg")
@@ -263,6 +261,9 @@ print(generated_texts)
 ```
 
 ### Now train MoE gating function
+
+We train the gating layers by providing sample images/prompts for each of the three experts. Here is a simple example training set: 
+
 ```python
 image_1 = Image.open("./VALIDATION/Q15.jpg") 
 image_1a = Image.open("./VALIDATION/Q31.jpg") 
@@ -296,9 +297,10 @@ moe_model.set_gating_layer_params(gating_layer_params)
 
 ![image/png](https://cdn-uploads.huggingface.co/production/uploads/623ce1c6b66fedf374859fe7/mh4eFDuFsTBOYbjc38PYz.png)
 
-Inference after MoE gating layers are trained:
 
-```
+Now that the MoE model has been trained, we can try inference.  Inference after MoE gating layers are trained:
+
+```python
 from transformers.image_utils import load_image
 
 image = load_image("https://d2r55xnwy6nx47.cloudfront.net/uploads/2018/02/Ants_Lede1300.jpg")
@@ -328,6 +330,8 @@ print(generated_texts)
 
 ### Push to hub and save locally
 
+We can save the MoE model either in Hugging Face Hub or locally:
+
 ```python
 repo_id='...'
 moe_name='Cephalo-Idefics2-3x8b-beta'
@@ -337,8 +341,22 @@ moe_model.push_to_hub (f'{repo_id}/'+merged_name, )
 ```
 
 Save locally:
-```
+```python
 processor.save_pretrained(moe_name, )
 moe_model.save_pretrained(moe_name,  )
 
+```
+
+Loading the model works as done above. Here included again for completeness:
+```python
+model_name_moe = f'{repo_id}/'+moe_name
+config = AutoConfig.from_pretrained(model_name_moe, trust_remote_code=True)
+processor = AutoProcessor.from_pretrained(model_name_moe, trust_remote_code=True) 
+moe_model = AutoModelForCausalLM.from_pretrained(
+    model_name_moe,config=config,
+    trust_remote_code=True,  torch_dtype=torch.bfloat16,   
+   # quantization_config=quantization_config,
+).to(device)
+
+count_parameters(moe_model)
 ```