set model and context from outside the model

LLM.py

@@ -33,18 +33,20 @@ class LLM:
     MODEL = "mosaicml/mpt-7b-chat"
     CONSTEXT = "You are an helpfull assistante in a school. You are helping a student with his homework."
 
-    def __init__(self):
-        self.load_model()
+    def __init__(self, model_name=None):
+        if model_name is None:
+            model_name = LLM.MODEL
+        self.load_model(model_name)
 
-    def load_model(self):
+    def load_model(self, model_name):
         device = "cuda:0" if torch.cuda.is_available() else "cpu"
-        tokenizer = AutoTokenizer.from_pretrained(LLM.MODEL)
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
 
         if device == "cuda:0":
-            model = AutoModelForCausalLM.from_pretrained(LLM.MODEL, trust_remote_code=True,
+            model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True,
                                                          torch_dtype=torch.float16, device_map="auto", load_in_8bit=True)
         else:
-            model = AutoModelForCausalLM.from_pretrained(LLM.MODEL, trust_remote_code=True)
+            model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
         pipe = pipeline(
             "text-generation",
             model=model,

app.py

@@ -1,7 +1,10 @@
 from LLM import LLM
 
-llm = LLM()
-chat = llm.get_chat(context=LLM.CONSTEXT)
+model = "mosaicml/mpt-7b-chat"
+context = "You are an helpfully assistant in a school. You are helping a student with his homework."
+
+llm = LLM(model)
+chat = llm.get_chat(context=context)
 while True:
     qn = input("Question: ")
     if qn == "exit":