Update app.py

app.py

@@ -1,7 +1,7 @@
 import optimum
 import transformers
 from transformers import AutoConfig, AutoTokenizer, AutoModel, AutoModelForCausalLM
-from optimum.bettertransformer import BetterTransformer
+from vllm import LLM, SamplingParams
 import torch
 import gradio as gr
 import json
@@ -19,10 +19,7 @@ repetition_penalty=1.7
 
 model_name = "Pclanglais/Bellay"
 
-tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
-model = transformers.AutoModelForCausalLM.from_pretrained(model_name,
-    device_map="auto"
-)
+llm = LLM(model_name)
 
 
 styles_prompts_dict = {
@@ -68,28 +65,13 @@ class MistralChatBot:
 
         system_prompt = styles_prompts_dict[style]
 
-        # Combine the user and assistant messages into a conversation
-        conversation = "<|im_start|>system\n" + system_prompt + "<|im_end|>\n<|im_start|>user\n" + user_message + "<|im_end|>\n<|im_start|>assistant\n"
-        # Encode the conversation using the tokenizer
-        input_ids = tokenizer.encode(conversation, return_tensors="pt", add_special_tokens=True)
-        input_ids = input_ids.to(device)
-        # Generate a response using the Falcon model
-        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=0.5, 
-            do_sample=True, 
-            max_new_tokens=max_new_tokens, 
-            top_p=top_p, 
-            repetition_penalty=repetition_penalty
-        )
-        # Decode the generated response to text
-        response_text = tokenizer.decode(response[0], skip_special_tokens=True)
-        return response_text
+        sampling_params = SamplingParams(temperature=0.7, top_p=.95, max_tokens=500, presence_penalty = 2)
+        detailed_prompt = "<|im_start|>system\n" + system_prompt + "<|im_end|>\n<|im_start|>user"""
+        detailed_prompt = detailed_prompt + "\n" + user_input + "<|im_end|>\n<|im_start|>assistant\n"
+        prompts = [detailed_prompt]
+        outputs = llm.generate(prompts, sampling_params, use_tqdm = False)
+        generated_text = outputs[0].outputs[0].text
+        return generated_text
 
     def predict_simple(self, user_message, style):
         system_prompt = styles_prompts_dict[style]