adapt zephyr miner to hf miner

neurons/miners/huggingface/miner.py

@@ -32,11 +32,10 @@
 
 class HuggingFaceMiner(Miner):
     """
-    Base miner which runs zephyr (https://huggingface.co/HuggingFaceH4/zephyr-7b-beta)    
-    This requires a GPU with at least 20GB of memory.
+    Base 🤗 Hugging Face miner, integrated with hf pipeline.    
     To run this miner from the project root directory:
 
-    python neurons/miners/zephyr/miner.py --wallet.name <wallet_name> --wallet.hotkey <wallet_hotkey> --subtensor.network <network> --netuid <netuid> --axon.port <port> --axon.external_port <port> --logging.debug True --neuron.model_id HuggingFaceH4/zephyr-7b-beta --neuron.system_prompt "Hello, I am a chatbot. I am here to help you with your questions." --neuron.max_tokens 64 --neuron.do_sample True --neuron.temperature 0.9 --neuron.top_k 50 --neuron.top_p 0.95 --wandb.on True --wandb.entity sn1 --wandb.project_name miners_experiments
+    python neurons/miners/huggingface/miner.py --wallet.name <wallet_name> --wallet.hotkey <wallet_hotkey> --neuron.model_id <model_id> --subtensor.network <network> --netuid <netuid> --axon.port <port> --axon.external_port <port> --logging.debug True --neuron.model_id HuggingFaceH4/zephyr-7b-beta --neuron.system_prompt "Hello, I am a chatbot. I am here to help you with your questions." --neuron.max_tokens 64 --neuron.do_sample True --neuron.temperature 0.9 --neuron.top_k 50 --neuron.top_p 0.95 --wandb.on True --wandb.entity sn1 --wandb.project_name miners_experiments
     """
     @classmethod
     def add_args(cls, parser: argparse.ArgumentParser):
@@ -49,27 +48,36 @@ def __init__(self, config=None):
         super().__init__(config=config)
 
         model_kwargs = None
-        if self.config.neuron.load_quantized:
-            bt.logging.info("Loading quantized model...")
+        if self.config.neuron.load_in_8bit:
+            bt.logging.info("Loading 8 bit quantized model...")
             model_kwargs = dict(
                 torch_dtype=torch.float16,
                 load_in_8bit=True,
             )
 
+        if self.config.neuron.load_in_4bit:
+            bt.logging.info("Loading 4 bit quantized model...")
+            model_kwargs = dict(
+                torch_dtype=torch.float32,
+                load_in_4bit=True,
+            )
+
         if self.config.wandb.on:
             self.identity_tags = ("hf_miner", )
 
-            if self.config.neuron.load_quantized:
-                self.identity_tags += ("8bits_quantization", )
+            if self.config.neuron.load_in_8bit:
+                self.identity_tags += ("8bit_quantization", )            
+            elif self.config.neuron.load_in_4bit:
+                self.identity_tags += ("4bit_quantization", )
 
         self.llm_pipeline = load_pipeline(
-            model_id=self.config.neuron.model_id,
-            torch_dtype=torch.float16,
+            model_id=self.config.neuron.model_id,            
             device=self.device,
             mock=self.config.mock,
             model_kwargs=model_kwargs,
         )        
 
+        self.model_id = self.config.neuron.model_id
         self.system_prompt = "You are a friendly chatbot who always responds concisely and helpfully. You are honest about things you don't know."
 
     async def forward(self, synapse: PromptingSynapse) -> PromptingSynapse:
@@ -92,7 +100,7 @@ async def forward(self, synapse: PromptingSynapse) -> PromptingSynapse:
             bt.logging.debug(f"📧 Message received, forwarding synapse: {synapse}")
 
             prompt = synapse.messages[-1]
-            bt.logging.debug(f"💬 Querying zephyr: {prompt}")
+            bt.logging.debug(f"💬 Querying {self.model_id}: {prompt}")
 
             response = HuggingFaceLLM(
                 llm_pipeline=self.llm_pipeline,