Update app.py

app.py

@@ -5,6 +5,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import gradio as gr
 import sentencepiece
+import gc
 
 title = "# Welcome to 🙋🏻‍♂️Tonic's🌷Tulu Chat!"
 description = """[allenai/tulu-2-dpo-7b](https://huggingface.co/allenai/tulu-2-dpo-7b) and larger Tulu-2 models are Instruct Llama Finetunes using the [mistralai/Mistral-7B](https://huggingface.co/mistralai/Mistral-7B-v0.1) recipe. You can use [allenai/tulu-2-13b](https://huggingface.co/allenai/tulu-2-13b) here via API using Gradio by scrolling down and clicking Use 'Via API' or privately by [cloning this space on huggingface](https://huggingface.co/spaces/Tonic1/TuluDemo?duplicate=true) See also the large model here : [allenai/tulu-2-dpo-70b](https://huggingface.co/allenai/tulu-2-dpo-70b) . [Join my active builders' server on discord](https://discord.gg/VqTxc76K3u). Let's build together!. [Add this Space as a discord bot to your server by clicking this link](https://discord.com/oauth2/authorize?client_id=1176628808212828231&scope=bot+applications.commands&permissions=326417525824). Big thanks to 🤗Huggingface Organisation for the🫂Community Grant"""
@@ -39,26 +40,31 @@ class TuluChatBot:
         return prompt
 
     def Tulu(self, user_message, temperature, max_new_tokens, top_p, repetition_penalty, do_sample):
-        prompt = self.format_prompt(user_message)
-        inputs = self.tokenizer(prompt, return_tensors='pt', add_special_tokens=True)
-        input_ids = inputs["input_ids"].to(self.model.device)
-        attention_mask = inputs["attention_mask"].to(self.model.device)
-
-
-        output_ids = self.model.generate(
-            input_ids,
-            attention_mask=attention_mask,
-            max_length=input_ids.shape[1] + max_new_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            repetition_penalty=repetition_penalty,
-            do_sample=do_sample
-        )
-
-        response = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
-        response = response.strip()
-        response = response.split("<|assistant|>\n")[-1]
-        return response
+        try:
+            prompt = self.format_prompt(user_message)
+            inputs = self.tokenizer(prompt, return_tensors='pt', add_special_tokens=True)
+            input_ids = inputs["input_ids"].to(self.model.device)
+            attention_mask = inputs["attention_mask"].to(self.model.device)
+
+
+            output_ids = self.model.generate(
+                input_ids,
+                attention_mask=attention_mask,
+                max_length=input_ids.shape[1] + max_new_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                repetition_penalty=repetition_penalty,
+                do_sample=do_sample
+            )
+
+            response = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
+            response = response.strip()
+            response = response.split("<|assistant|>\n")[-1]
+            return response        
+        finally:
+            del input_ids, attention_mask, output_ids
+            gc.collect()
+            torch.cuda.empty_cache()
 
 def gradio_Tulu(user_message, system_message, max_new_tokens, temperature, top_p, repetition_penalty, do_sample):
     Tulu_bot.set_system_message(system_message)