Update app.py

app.py

@@ -71,9 +71,14 @@ class DocumentRetrievalAndGeneration:
         return generate_text
     def initialize_llm2(self,model_id):
                 
-        self.tokenizer = BloomTokenizerFast.from_pretrained("bigscience/bloom-3b")#, local_files_only=True)
-        self.model = BloomForCausalLM.from_pretrained("bigscience/bloom-3b")#, local_files_only=True)
-        self.result_length = 2048
+        model_name = "mistralai/Mistral-7B-Instruct-v0.2"
+        pipeline = transformers.pipeline(
+            "text-generation",
+            model=model_name,
+            model_kwargs={"torch_dtype": torch.bfloat16},
+            device="cpu",
+        )
+
         
         # return generate_text
 
@@ -150,9 +155,44 @@ class DocumentRetrievalAndGeneration:
 
         # decoded = self.llm.tokenizer.batch_decode(generated_ids)
         # generated_response = decoded[0]
-        inputs = self.tokenizer(prompt, return_tensors="pt")
-        generated_response=self.tokenizer.decode(self.model.generate(inputs["input_ids"], max_length=self.result_length,no_repeat_ngram_size=2)[0])
-        print(generated_response)
+        messages = []
+        # Check if history is None or empty and handle accordingly
+        if history:
+            for user_msg, assistant_msg in history:
+                messages.append({"role": "user", "content": user_msg})
+                messages.append({"role": "assistant", "content": assistant_msg})
+        
+        # Always add the current user message
+        messages.append({"role": "user", "content": message})
+        
+        # Construct the prompt using the pipeline's tokenizer
+        prompt = pipeline.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+    
+        # Generate the response
+        terminators = [
+            pipeline.tokenizer.eos_token_id,
+            pipeline.tokenizer.convert_tokens_to_ids("")
+        ]
+    
+        # Adjust the temperature slightly above given to ensure variety
+        adjusted_temp = temperature + 0.1
+    
+        # Generate outputs with adjusted parameters
+        outputs = pipeline(
+            prompt,
+            max_new_tokens=max_new_tokens,
+            do_sample=True,
+            temperature=adjusted_temp,
+            top_p=0.9
+        )
+    
+        # Extract the generated text, skipping the length of the prompt
+        generated_text = outputs[0]["generated_text"]
+        generated_response = generated_text[len(prompt):]
         
         match1 = re.search(r'\[/INST\](.*?)</s>', generated_response, re.DOTALL)