Update app.py

app.py

@@ -50,32 +50,26 @@ model_configs = [
 class ModelManager:
     def __init__(self):
         self.loaded = False
+        self.models = {}
 
     def load_model(self, model_config):
-        try:
-            return {"model": Llama.from_pretrained(repo_id=model_config['repo_id'], filename=model_config['filename']), "name": model_config['name']}
-        except Exception as e:
-            print(f"Error loading model {model_config['name']}: {e}")
-            return None
+        if model_config['name'] not in self.models:
+            try:
+                self.models[model_config['name']] = Llama.from_pretrained(repo_id=model_config['repo_id'], filename=model_config['filename'])
+            except Exception as e:
+                print(f"Error loading model {model_config['name']}: {e}")
 
     def load_all_models(self):
-        if self.loaded:
-            return global_data['models']
+        if not self.loaded:
+            with ThreadPoolExecutor() as executor:
+                for config in model_configs:
+                    executor.submit(self.load_model, config)
+            self.loaded = True
 
-        with ThreadPoolExecutor() as executor:
-            futures = [executor.submit(self.load_model, config) for config in model_configs]
-            models = []
-            for future in as_completed(futures):
-                model = future.result()
-                if model:
-                    models.append(model)
-
-        global_data['models'] = {model['name']: model['model'] for model in models}
-        self.loaded = True
-        return global_data['models']
+        return self.models
 
 model_manager = ModelManager()
-model_manager.load_all_models()
+global_data['models'] = model_manager.load_all_models()
 
 class ChatRequest(BaseModel):
     message: str
@@ -103,7 +97,7 @@ def remove_duplicates(text):
 def generate_model_response(model, inputs, top_k, top_p, temperature):
     try:
         response = model(inputs, top_k=top_k, top_p=top_p, temperature=temperature)
-        return remove_duplicates(response)
+        return remove_duplicates(response['choices'][0]['text'])
     except Exception as e:
         print(f"Error generating model response: {e}")
         return ""