Update app.py

app.py

@@ -44,28 +44,34 @@ model = transformers.AutoModelForCausalLM.from_pretrained(model_name,
 model.eval()
 
 class StarlingBot:
-    def __init__(self, system_prompt="The following dialogue is a conversation"):
-        self.system_prompt = system_prompt
+        def __init__(self, system_prompt="The following dialogue is a conversation"):
+            self.system_prompt = system_prompt
 
-    def predict(self, user_message, assistant_message, system_prompt, do_sample, temperature=0.4, max_new_tokens=700, top_p=0.99, repetition_penalty=1.9):
-        conversation = f" <s> [INST] {self.system_prompt} [INST]  {assistant_message if assistant_message else ''} </s> [/INST]  {user_message}  </s> "
-        input_ids = tokenizer.encode(conversation, return_tensors="pt", add_special_tokens=False)
-        input_ids = input_ids.to(device)
-        response = model.generate(
-            input_ids=input_ids, 
-            use_cache=False, 
-            early_stopping=False, 
-            bos_token_id=model.config.bos_token_id, 
-            eos_token_id=model.config.eos_token_id, 
-            pad_token_id=model.config.eos_token_id, 
-            temperature=temperature, 
-            do_sample=True, 
-            max_new_tokens=max_new_tokens, 
-            top_p=top_p, 
-            repetition_penalty=repetition_penalty
-        )
-        response_text = tokenizer.decode(response[0], skip_special_tokens=True)
-        return response_text
+        def predict(self, user_message, assistant_message, system_prompt, do_sample, temperature=0.4, max_new_tokens=700, top_p=0.99, repetition_penalty=1.9):
+            conversation = f" <s> [INST] {self.system_prompt} [INST]  {assistant_message if assistant_message else ''} </s> [/INST]  {user_message}  </s> "
+            input_ids = tokenizer.encode(conversation, return_tensors="pt", add_special_tokens=False)
+            input_ids = input_ids.to(device)
+            response = model.generate(
+                input_ids=input_ids, 
+                use_cache=False, 
+                early_stopping=False, 
+                bos_token_id=model.config.bos_token_id, 
+                    eos_token_id=model.config.eos_token_id, 
+                pad_token_id=model.config.eos_token_id, 
+                temperature=temperature, 
+                do_sample=True, 
+                max_new_tokens=max_new_tokens, 
+                top_p=top_p, 
+                repetition_penalty=repetition_penalty
+            )
+            response_text = tokenizer.decode(response[0], skip_special_tokens=True)
+            response_text = response.strip()
+#           response_text = response.split("<|assistant|>\n")[-1]
+            return response_text       
+        finally:
+            del input_ids, attention_mask, output_ids
+            gc.collect()
+            torch.cuda.empty_cache()
 
 starling_bot = StarlingBot()