Added Client class

shared.py

@@ -2,6 +2,10 @@ import yaml
 
 from typing import Dict
 from pydantic import BaseModel, ValidationError
+from huggingface_hub import hf_hub_download
+from huggingface_hub.utils import EntryNotFoundError
+
+from openai import OpenAI
 
 
 
@@ -30,4 +34,46 @@ class ModelConfig(BaseModel):
         try:
             return cls(**data)
         except ValidationError as e:
-            raise e
+            raise e
+        
+
+class Client:
+    def __init__(self, api_url, api_key):
+        self.api_url = api_url
+        self.api_key = api_key
+        
+        self.init_all()
+
+    def init_all(self):
+        self.init_client()
+        self.get_metadata()
+        self.get_personas()
+
+    def init_client(self):
+        self.client = OpenAI(
+            base_url=f"{self.api_url}/v1",
+            api_key=self.api_key,
+        )
+
+    def get_metadata(self):
+        models = self.client.models.list()
+        vllm_model_name = models.data[0].id
+
+        model_name, *suffix = vllm_model_name.split("@")
+        revision = dict(enumerate(suffix)).get(0, None)
+
+        self.vllm_model_name = vllm_model_name
+        self.model_name = model_name
+        self.revision = revision
+    
+    def get_personas(self):
+        try:
+            config_path = hf_hub_download(self.model_name, "config.yaml",
+                                    subfolder="datasets",
+                                    revision=self.revision)
+            self.config = ModelConfig.from_yaml(config_path)
+            self.personas = self.config.pile.persona2system
+        except EntryNotFoundError:
+            self.personas = {}
+
+        self.personas["default"] = None