Update app.py

app.py

@@ -26,13 +26,25 @@ gcp_credentials = json.loads(GOOGLE_CLOUD_CREDENTIALS)
 storage_client = storage.Client.from_service_account_info(gcp_credentials)
 bucket = storage_client.bucket(GOOGLE_CLOUD_BUCKET)
 
+MODEL_NAMES = {
+    "starcoder": "starcoder2-3b-q2_k.gguf",
+    "gemma_2b_it": "gemma-2-2b-it-q2_k.gguf",
+    "llama_3_2_1b": "Llama-3.2-1B.Q2_K.gguf",
+    "gemma_2b_imat": "gemma-2-2b-iq1_s-imat.gguf",
+    "phi_3_mini": "phi-3-mini-128k-instruct-iq2_xxs-imat.gguf",
+    "qwen2_0_5b": "qwen2-0.5b-iq1_s-imat.gguf",
+    "gemma_9b_it": "gemma-2-9b-it-q2_k.gguf",
+    "gpt2_xl": "gpt2-xl-q2_k.gguf",
+}
+
 class ModelManager:
     def __init__(self):
         self.params = {"n_ctx": 2048, "n_batch": 512, "n_predict": 512, "repeat_penalty": 1.1, "n_threads": 1, "seed": -1, "stop": ["</s>"], "tokens": []}
-#        self.tokenizer = LlamaTokenizer.from_pretrained("decapoda-research/llama-7b-hf") #Load tokenizer from GCS for production
+#        self.tokenizer = LlamaTokenizer.from_pretrained("decapoda-research/llama-7b-hf") # Load from GCS for production
         self.request_queue = Queue()
         self.response_queue = Queue()
-        self.model = self.load_model_from_bucket("llama-2-7b-chat/llama-2-7b-chat.Q4_K_M.gguf")
+        self.models = {} # Dictionary to hold multiple models
+        self.load_models()
         self.start_processing_processes()
 
     def load_model_from_bucket(self, bucket_path):
@@ -44,6 +56,12 @@ class ModelManager:
             print(f"Error loading model: {e}")
             return None
 
+    def load_models(self):
+        for name, path in MODEL_NAMES.items():
+            model = self.load_model_from_bucket(path)
+            if model:
+                self.models[name] = model
+
     def save_model_to_bucket(self, model, bucket_path):
         blob = bucket.blob(bucket_path)
         try:
@@ -72,14 +90,15 @@ class ModelManager:
             print(f"Error during training: {e}")
 
 
-    def generate_text(self, prompt):
-        if self.model:
+    def generate_text(self, prompt, model_name):
+        if model_name in self.models:
+            model = self.models[model_name]
             inputs = self.tokenizer(prompt, return_tensors="pt")
-            outputs = self.model.generate(**inputs, max_new_tokens=100)
+            outputs = model.generate(**inputs, max_new_tokens=100)
             generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
             return generated_text
         else:
-            return "Error loading model."
+            return "Error: Model not found."
 
     def start_processing_processes(self):
         p = Process(target=self.process_requests)
@@ -90,9 +109,9 @@ class ModelManager:
             request_data = self.request_queue.get()
             if request_data is None:
                 break
-            inputs, top_p, top_k, temperature, max_tokens = request_data
+            inputs, model_name, top_p, top_k, temperature, max_tokens = request_data
             try:
-                response = self.generate_text(inputs)
+                response = self.generate_text(inputs, model_name)
                 self.response_queue.put(response)
             except Exception as e:
                 print(f"Error during inference: {e}")
@@ -102,30 +121,34 @@ model_manager = ModelManager()
 
 class ChatRequest(BaseModel):
     message: str
+    model_name: str
 
 @spaces.GPU()
-async def generate_streaming_response(inputs):
+async def generate_streaming_response(inputs, model_name):
     top_p = 0.9
     top_k = 50
     temperature = 0.7
     max_tokens = model_manager.params["n_ctx"] - len(model_manager.tokenizer.encode(inputs))
-    model_manager.request_queue.put((inputs, top_p, top_k, temperature, max_tokens))
+    model_manager.request_queue.put((inputs, model_name, top_p, top_k, temperature, max_tokens))
     full_text = model_manager.response_queue.get()
     async def stream_response():
         yield full_text
     return StreamingResponse(stream_response())
 
-async def process_message(message):
+async def process_message(message, model_name):
     inputs = message.strip()
-    return await generate_streaming_response(inputs)
+    return await generate_streaming_response(inputs, model_name)
 
 @app.post("/generate_multimodel")
 async def api_generate_multimodel(request: Request):
     data = await request.json()
     message = data["message"]
-    return await process_message(message)
+    model_name = data.get("model_name", list(MODEL_NAMES.keys())[0]) 
+    if model_name not in MODEL_NAMES:
+        return {"error": "Invalid model name"}
+    return await process_message(message, model_name)
 
-iface = gr.Interface(fn=process_message, inputs=gr.Textbox(lines=2, placeholder="Enter your message here..."), outputs=gr.Markdown(stream=True), title="Unified Multi-Model API", description="Enter a message to get responses from a unified model.") #gradio is not suitable for production
+iface = gr.Interface(fn=process_message, inputs=[gr.Textbox(lines=2, placeholder="Enter your message here..."), gr.Dropdown(list(MODEL_NAMES.keys()), label="Select Model")], outputs=gr.Markdown(stream=True), title="Unified Multi-Model API", description="Enter a message to get responses from a unified model.") #gradio is not suitable for production
 
 if __name__ == "__main__":
     iface.launch()