Update app.py

app.py

@@ -2,134 +2,134 @@ import gradio as gr
 from gradio_calendar import Calendar
 
 
-# # from langchain.embeddings.huggingface import HuggingFaceEmbeddings
-# # from langchain_community.vectorstores import Chroma
-# # from langchain_core.output_parsers import StrOutputParser
-# import torch
-# from transformers import (
-#     AutoModelForCausalLM,
-#     AutoTokenizer,
-#     BitsAndBytesConfig,
-#     pipeline,
-#     StoppingCriteria, StoppingCriteriaList
-# )
 
-# from langchain.prompts import PromptTemplate
-# from langchain_community.llms import HuggingFacePipeline
-# from langchain.chains import LLMChain
-# from langchain_core.runnables import RunnablePassthrough, RunnableParallel
+# from langchain.embeddings.huggingface import HuggingFaceEmbeddings
+# from langchain_community.vectorstores import Chroma
+# from langchain_core.output_parsers import StrOutputParser
+import torch
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    BitsAndBytesConfig,
+    pipeline,
+    StoppingCriteria, StoppingCriteriaList
+)
 
+from langchain.prompts import PromptTemplate
+from langchain_community.llms import HuggingFacePipeline
+from langchain.chains import LLMChain
+from langchain_core.runnables import RunnablePassthrough, RunnableParallel
 
-# # instructor_embeddings = HuggingFaceEmbeddings(model_name="intfloat/multilingual-e5-large",
-# #                                                       model_kwargs={"device": "cuda"})
 
+instructor_embeddings = HuggingFaceEmbeddings(model_name="intfloat/multilingual-e5-large",
+                                                      model_kwargs={"device": "cuda"})
 
 
-# model_name='SherlockAssistant/Mistral-7B-Instruct-Ukrainian'
+model_name='SherlockAssistant/Mistral-7B-Instruct-Ukrainian'
 
-# tokenizer = AutoTokenizer.from_pretrained(model_name)
-# #tokenizer.pad_token = tokenizer.unk_token
-# #tokenizer.padding_side = "right"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+#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"
+# # Compute dtype for 4-bit base models
+# bnb_4bit_compute_dtype = "float16"
 
-# # # Quantization type (fp4 or nf4)
-# # bnb_4bit_quant_type = "nf4"
+# # Quantization type (fp4 or nf4)
+# bnb_4bit_quant_type = "nf4"
 
-# # # Activate nested quantization for 4-bit base models (double quantization)
-# # use_nested_quant = False
+# # 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)
+# #################################################################
+# # 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,
-# # )
+# 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)
+# # 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)
-# 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]
+model = AutoModelForCausalLM.from_pretrained(
+    model_name)
+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
+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()])
-
-
-# text_generation_pipeline = pipeline(
-#     model=model,
-#     tokenizer=tokenizer,
-#     task="text-generation",
-#     temperature=0.01,
-#     repetition_penalty=1.2,
-#     return_full_text=True,
-#     max_new_tokens=750, do_sample=True,
-#     top_k=50, top_p=0.95,
-#     stopping_criteria=stopping_criteria
-# )
-# mistral_llm = HuggingFacePipeline(pipeline=text_generation_pipeline)
-# # # # load chroma from disk
-# db3 = Chroma(persist_directory="/chroma/", embedding_function=instructor_embeddings)
-
-
-
-
-# retriever = db3.as_retriever(search_type="similarity_score_threshold",
-#                                  search_kwargs={"score_threshold": .5,
-#                                                 "k": 20})
-
-# #retriever = db3.as_retriever(search_kwargs={"k":15})
-# # Get pre-written rag prompt
-# def format_docs(docs):
-#    return "\n\n".join(doc.page_content for doc in docs)
-
-
-# template ="""" [INST] Ти асистент для надання відповідей з законодавства України. Використовуй лише вказаний нижче Context максимально точно. Описуй лише події простими словами без формальностей. Пиши чотири речення і будь максимально точним.  Якщо контекст пустий - пиши "Я не маю релевантної інформації. Спробуйте ще".
-# Context: {context}
-# ### QUESTION:
-# {question}
-# [/INST]
-# """
-# prompt = PromptTemplate(
-#     input_variables=["context", "question"],
-#     template=template,
-# )
-
-# rag_chain_from_docs = (
-#     RunnablePassthrough.assign(context=(lambda x: format_docs(x["context"])))
-#     | prompt
-#     | mistral_llm
-#     | StrOutputParser()
-# )
-
-# rag_chain_with_source = RunnableParallel(
-#     {"context": retriever, "question": RunnablePassthrough()}
-# ).assign(answer=rag_chain_from_docs)
+stopping_criteria = StoppingCriteriaList([StopOnTokens()])
+
+
+text_generation_pipeline = pipeline(
+    model=model,
+    tokenizer=tokenizer,
+    task="text-generation",
+    temperature=0.01,
+    repetition_penalty=1.2,
+    return_full_text=True,
+    max_new_tokens=750, do_sample=True,
+    top_k=50, top_p=0.95,
+    stopping_criteria=stopping_criteria
+)
+mistral_llm = HuggingFacePipeline(pipeline=text_generation_pipeline)
+# # # load chroma from disk
+db3 = Chroma(persist_directory="/chroma/", embedding_function=instructor_embeddings)
+
+
+
+
+retriever = db3.as_retriever(search_type="similarity_score_threshold",
+                                 search_kwargs={"score_threshold": .5,
+                                                "k": 20})
+
+#retriever = db3.as_retriever(search_kwargs={"k":15})
+# Get pre-written rag prompt
+def format_docs(docs):
+   return "\n\n".join(doc.page_content for doc in docs)
+
+
+template ="""" [INST] Ти асистент для надання відповідей з законодавства України. Використовуй лише вказаний нижче Context максимально точно. Описуй лише події простими словами без формальностей. Пиши чотири речення і будь максимально точним.  Якщо контекст пустий - пиши "Я не маю релевантної інформації. Спробуйте ще".
+Context: {context}
+### QUESTION:
+{question}
+[/INST]
+"""
+prompt = PromptTemplate(
+    input_variables=["context", "question"],
+    template=template,
+)
+
+rag_chain_from_docs = (
+    RunnablePassthrough.assign(context=(lambda x: format_docs(x["context"])))
+    | prompt
+    | mistral_llm
+    | StrOutputParser()
+)
+
+rag_chain_with_source = RunnableParallel(
+    {"context": retriever, "question": RunnablePassthrough()}
+).assign(answer=rag_chain_from_docs)