Update app.py

app.py

@@ -11,7 +11,8 @@ from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
     BitsAndBytesConfig,
-    pipeline
+    pipeline,
+    StoppingCriteria, StoppingCriteriaList
 )
 
 
@@ -38,45 +39,60 @@ instructor_embeddings = HuggingFaceEmbeddings(model_name="intfloat/multilingual-
 model_name='SherlockAssistant/Mistral-7B-Instruct-Ukrainian'
 
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-tokenizer.pad_token = tokenizer.eos_token
+tokenizer.pad_token = tokenizer.unk_token
 tokenizer.padding_side = "right"
 
 
-# Activate 4-bit precision base model loading
-use_4bit = True
+# # Activate 4-bit precision base model loading
+# use_4bit = True
+
+# # Compute dtype for 4-bit base models
+# bnb_4bit_compute_dtype = "float16"
+
+# # Quantization type (fp4 or nf4)
+# bnb_4bit_quant_type = "nf4"
+
+# # Activate nested quantization for 4-bit base models (double quantization)
+# use_nested_quant = False
+
+# #################################################################
+# # Set up quantization config
+# #################################################################
+# compute_dtype = getattr(torch, bnb_4bit_compute_dtype)
+
+# bnb_config = BitsAndBytesConfig(
+#     load_in_4bit=use_4bit,
+#     bnb_4bit_quant_type=bnb_4bit_quant_type,
+#     bnb_4bit_compute_dtype=compute_dtype,
+#     bnb_4bit_use_double_quant=use_nested_quant,
+# )
+
+# # Check GPU compatibility with bfloat16
+# if compute_dtype == torch.float16 and use_4bit:
+#     major, _ = torch.cuda.get_device_capability()
+#     if major >= 8:
+#         print("=" * 80)
+#         print("Your GPU supports bfloat16: accelerate training with bf16=True")
+#         print("=" * 80)
+
+
+# model = AutoModelForCausalLM.from_pretrained(
+#     model_name,
+#     quantization_config=bnb_config,
+# )
+stop_list = [" \n\nAnswer:", " \n", " \n\n"]
+stop_token_ids = [tokenizer(x,  return_tensors='pt', add_special_tokens=False)['input_ids'] for x in stop_list]
+stop_token_ids = [torch.LongTensor(x).to("cuda") for x in stop_token_ids]
+    
+class StopOnTokens(StoppingCriteria):
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        for stop_ids in stop_token_ids:
+            if torch.eq(input_ids[0][-len(stop_ids[0])+1:], stop_ids[0][1:]).all():
+                return True
+        return False
+        
+stopping_criteria = StoppingCriteriaList([StopOnTokens()])
 
-# Compute dtype for 4-bit base models
-bnb_4bit_compute_dtype = "float16"
-
-# Quantization type (fp4 or nf4)
-bnb_4bit_quant_type = "nf4"
-
-# Activate nested quantization for 4-bit base models (double quantization)
-use_nested_quant = False
-
-
-compute_dtype = getattr(torch, bnb_4bit_compute_dtype)
-
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=use_4bit,
-    bnb_4bit_quant_type=bnb_4bit_quant_type,
-    bnb_4bit_compute_dtype=compute_dtype,
-    bnb_4bit_use_double_quant=use_nested_quant,
-)
-
-# Check GPU compatibility with bfloat16
-if compute_dtype == torch.float16 and use_4bit:
-    major, _ = torch.cuda.get_device_capability()
-    if major >= 8:
-        print("=" * 80)
-        print("Your GPU supports bfloat16: accelerate training with bf16=True")
-        print("=" * 80)
-
-
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    quantization_config=bnb_config,
-)
 
 text_generation_pipeline = pipeline(
     model=model,
@@ -86,10 +102,11 @@ text_generation_pipeline = pipeline(
     repetition_penalty=1.2,
     return_full_text=True,
     max_new_tokens=750, do_sample=True,
-    top_k=50, top_p=0.95
+    top_k=50, top_p=0.95,
+    stopping_criteria=stopping_criteria
 )
 mistral_llm = HuggingFacePipeline(pipeline=text_generation_pipeline)
-# # load chroma from disk
+# # # load chroma from disk
 db3 = Chroma(persist_directory="/content/gdrive/MyDrive/diploma/all_emb/chroma/", embedding_function=instructor_embeddings)
 
 
@@ -130,13 +147,11 @@ rag_chain_with_source = RunnableParallel(
 
 
 def format_result(result):
-    # Extract unique pairs of titles and video IDs from the context
     unique_videos = set((doc.metadata['title'], doc.metadata['act_url']) for doc in result['context'])
 
     # Create a plain text string where each title is followed by its URL
     titles_with_links = [
-    unique_videos = set((doc.metadata['title'], doc.metadata['act_url']) for doc in result['context'])
-        f"{title}: {act_url}" for title,  act_url in unique_videos
+        f"{title}: {act_url}" for title, act_url in unique_videos
     ]
 
     # Join these entries with line breaks to form a clear list
@@ -154,7 +169,7 @@ def format_result(result):
 def generate_with_filters(message, subject_input, rubric, date_beg, date_end):
     if date_beg == "2010-01-01" and date_end == "2025-01-01":
         rag_chain_with_filters = RunnableParallel(
-             {"context": db3.as_retriever(search_type="mmr", search_kwargs={"k": 15,
+             {"context": db3.as_retriever(search_type="mmr", search_kwargs={"k": 10,
                                                                                        "filter": {'$and': [{'subject': {
                                                                                            '$in': subject_input}}, {
                                                                                            'rubric': {
@@ -163,11 +178,11 @@ def generate_with_filters(message, subject_input, rubric, date_beg, date_end):
         ).assign(answer=rag_chain_from_docs)
     else:
         rag_chain_with_filters = RunnableParallel(
-        {"context": db3.as_retriever(search_type="mmr", search_kwargs={"k": 15,
+        {"context": db3.as_retriever(search_type="mmr", search_kwargs={"k": 10,
                                                                            "filter": {'$and': [{'subject': {
                                                                                '$in': subject_input}}, {
                                                                                                    'rubric': {
-                                                                                                       '$in': rubric}},"act_date": {"$gte": date_beg}, "act_date": {"$lte": date_end}] }}), "question": RunnablePassthrough()}
+                                                                                                       '$in': rubric}},{"act_date": {"$gte": date_beg}}, {"act_date": {"$lte": date_end}}] }}), "question": RunnablePassthrough()}
     ).assign(answer=rag_chain_from_docs)
     result = rag_chain_with_filters.invoke(message)
     return result