bug fix for methods

app.py

@@ -18,22 +18,20 @@ tokenizer = None
 
 def get_model_and_tokenizer(model_id):
     global model, tokenizer
-    try:
-        print(f"Loading tokenizer for model_id: {model_id}")
-        # Load the tokenizer
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
-        tokenizer.pad_token = tokenizer.eos_token
-        
-        print(f"Loading model and for model_id: {model_id}")
-        # Load the model
-        model = AutoModelForCausalLM.from_pretrained(model_id) #, device_map="auto")
-        model.config.use_cache = False
-   
-    except Exception as e:
-        print(f"Error loading model: {e}")
- 
-    return "No complete blocks found. Please check the format of the response."
- 
+    if model is None or tokenizer is None:
+        try:
+            print(f"Loading tokenizer for model_id: {model_id}")
+            tokenizer = AutoTokenizer.from_pretrained(model_id)
+            tokenizer.pad_token = tokenizer.eos_token
+            
+            print(f"Loading model for model_id: {model_id}")
+            model = AutoModelForCausalLM.from_pretrained(model_id)
+            model.config.use_cache = False
+        except Exception as e:
+            print(f"Error loading model: {e}")
+            raise e  # Raise the error to be caught in the POST request
+    else:
+        print(f"Model and tokenizer for {model_id} are already loaded.")
 
         # max_new_tokens=100,
         # min_length=5,
@@ -70,7 +68,10 @@ def get_model_and_tokenizer(model_id):
         #truncation=True,              # Enable truncation for longer prompts
        #
 
-def generate_response(user_input):
+def generate_response(user_input, model_id):
+    # Ensure model and tokenizer are loaded
+    get_model_and_tokenizer(model_id)  # Load the model/tokenizer if not already loaded
+    
     prompt = formatted_prompt(user_input)
     inputs = tokenizer([prompt], return_tensors="pt")
 
@@ -87,7 +88,7 @@ def generate_response(user_input):
 
     outputs = model.generate(**inputs, generation_config=generation_config)
     response = tokenizer.decode(outputs[:, inputs['input_ids'].shape[-1]:][0], skip_special_tokens=True)
-    return response.strip().split("Assistant:")[-1].strip()  # Get the part after 'Assistant:'
+    return response.strip().split("Assistant:")[-1].strip()
      
 def formatted_prompt(question) -> str:
     return f"<|startoftext|>User: {question}\nAssistant:"