Update app.py

app.py

@@ -7,10 +7,12 @@ from langchain.vectorstores import FAISS
 from langchain.chains import ConversationalRetrievalChain
 import gradio as gr
 from langchain.embeddings import HuggingFaceEmbeddings
-from sentence_transformers import CrossEncoder
 
+
+# Load the Hugging Face token from environment
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 
+# Define stopping criteria
 class StopOnTokens(StoppingCriteria):
     def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
         for stop_ids in stop_token_ids:
@@ -18,9 +20,11 @@ class StopOnTokens(StoppingCriteria):
                 return True
         return False
 
+# Load the LLaMA model and tokenizer
 model_id = 'meta-llama/Meta-Llama-3-8B-Instruct'
 device = f'cuda:{cuda.current_device()}' if cuda.is_available() else 'cpu'
 
+# Set quantization configuration
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_quant_type='nf4',
@@ -31,11 +35,13 @@ bnb_config = BitsAndBytesConfig(
 tokenizer = AutoTokenizer.from_pretrained(model_id, token=HF_TOKEN)
 model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", token=HF_TOKEN, quantization_config=bnb_config)
 
+# Define stopping criteria
 stop_list = ['\nHuman:', '\n```\n']
 stop_token_ids = [tokenizer(x)['input_ids'] for x in stop_list]
 stop_token_ids = [torch.LongTensor(x).to(device) for x in stop_token_ids]
 stopping_criteria = StoppingCriteriaList([StopOnTokens()])
 
+# Create text generation pipeline
 generate_text = pipeline(
     model=model,
     tokenizer=tokenizer,
@@ -49,21 +55,19 @@ generate_text = pipeline(
 
 llm = HuggingFacePipeline(pipeline=generate_text)
 
-"""Load the stored FAISS index"""
+# Load the stored FAISS index
 try:
-    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2", model_kwargs={"device": "cuda"})
-    vectorstore = FAISS.load_local('faiss_index', embeddings)
-    print("Loaded embeddings from FAISS Index successfully")
+    vectorstore = FAISS.load_local('faiss_index', HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2", model_kwargs={"device": "cuda"}))
+    print("Loaded embedding successfully")
 except ImportError as e:
     print("FAISS could not be imported. Make sure FAISS is installed correctly.")
     raise e
 
+# Set up the Conversational Retrieval Chain
 chain = ConversationalRetrievalChain.from_llm(llm, vectorstore.as_retriever(), return_source_documents=True)
 
 chat_history = []
 
-reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
-
 def format_prompt(query):
     prompt = f"""
     You are a knowledgeable assistant with access to a comprehensive database. 
@@ -84,19 +88,8 @@ def format_prompt(query):
 
 def qa_infer(query):
     formatted_prompt = format_prompt(query)
-    results = chain({"question": formatted_prompt, "chat_history": chat_history})
-    
-    documents = results['source_documents']
-    query_document_pairs = [[query, doc.page_content] for doc in documents]
-    scores = reranker.predict(query_document_pairs)
-    
-    """Sort documents based on the re-ranker scores"""
-    ranked_docs = sorted(zip(scores, documents), key=lambda x: x[0], reverse=True)
-    
-    """Extract the best document"""
-    best_doc = ranked_docs[0][1].page_content if ranked_docs else ""
-    
-    return best_doc
+    result = chain({"question": formatted_prompt, "chat_history": chat_history})
+    return result['answer']
 
 EXAMPLES = ["How to use IPU1_0 instead of A15_0 to process NDK in TDA2x-EVM", 
             "Can BQ25896 support I2C interface?", 
@@ -104,3 +97,103 @@ EXAMPLES = ["How to use IPU1_0 instead of A15_0 to process NDK in TDA2x-EVM",
 
 demo = gr.Interface(fn=qa_infer, inputs="text", allow_flagging='never', examples=EXAMPLES, cache_examples=False, outputs="text")
 demo.launch()
+
+# import os
+# import torch
+# from torch import cuda, bfloat16
+# from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, BitsAndBytesConfig, StoppingCriteria, StoppingCriteriaList
+# from langchain.llms import HuggingFacePipeline
+# from langchain.vectorstores import FAISS
+# from langchain.chains import ConversationalRetrievalChain
+# import gradio as gr
+# from langchain.embeddings import HuggingFaceEmbeddings
+
+# # Load the Hugging Face token from environment
+# HF_TOKEN = os.environ.get("HF_TOKEN", None)
+
+# # Define stopping criteria
+# 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):], stop_ids).all():
+#                 return True
+#         return False
+
+# # Load the LLaMA model and tokenizer
+# model_id = 'meta-llama/Meta-Llama-3-8B-Instruct'
+# device = f'cuda:{cuda.current_device()}' if cuda.is_available() else 'cpu'
+
+# # Set quantization configuration
+# bnb_config = BitsAndBytesConfig(
+#     load_in_4bit=True,
+#     bnb_4bit_quant_type='nf4',
+#     bnb_4bit_use_double_quant=True,
+#     bnb_4bit_compute_dtype=bfloat16
+# )
+
+# tokenizer = AutoTokenizer.from_pretrained(model_id, token=HF_TOKEN)
+# model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", token=HF_TOKEN, quantization_config=bnb_config)
+
+# # Define stopping criteria
+# stop_list = ['\nHuman:', '\n```\n']
+# stop_token_ids = [tokenizer(x)['input_ids'] for x in stop_list]
+# stop_token_ids = [torch.LongTensor(x).to(device) for x in stop_token_ids]
+# stopping_criteria = StoppingCriteriaList([StopOnTokens()])
+
+# # Create text generation pipeline
+# generate_text = pipeline(
+#     model=model,
+#     tokenizer=tokenizer,
+#     return_full_text=True,
+#     task='text-generation',
+#     stopping_criteria=stopping_criteria,
+#     temperature=0.1,
+#     max_new_tokens=512,
+#     repetition_penalty=1.1
+# )
+
+# llm = HuggingFacePipeline(pipeline=generate_text)
+
+# # Load the stored FAISS index
+# try:
+#     embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2", model_kwargs={"device": "cuda"})
+#     vectorstore = FAISS.load_local('faiss_index', embeddings)
+#     print("Loaded embedding successfully")
+# except ImportError as e:
+#     print("FAISS could not be imported. Make sure FAISS is installed correctly.")
+#     raise e
+
+# # Set up the Conversational Retrieval Chain
+# chain = ConversationalRetrievalChain.from_llm(llm, vectorstore.as_retriever(), return_source_documents=True)
+
+# chat_history = []
+
+# def format_prompt(query):
+#     prompt = f"""
+#     You are a knowledgeable assistant with access to a comprehensive database. 
+#     I need you to answer my question and provide related information in a specific format.
+#     Here's what I need:
+#     1. A brief, general response to my question based on related answers retrieved.
+#     2. A JSON-formatted output containing:
+#        - "question": The original question.
+#        - "answer": The detailed answer.
+#        - "related_questions": A list of related questions and their answers, each as a dictionary with the keys:
+#          - "question": The related question.
+#          - "answer": The related answer.
+#     Here's my question:
+#     {query}
+#     Include a brief final answer without additional comments, sign-offs, or extra phrases. Be direct and to the point.
+#     """
+#     return prompt
+
+# def qa_infer(query):
+#     formatted_prompt = format_prompt(query)
+#     result = chain({"question": formatted_prompt, "chat_history": chat_history})
+#     return result['answer']
+
+# EXAMPLES = ["How to use IPU1_0 instead of A15_0 to process NDK in TDA2x-EVM", 
+#             "Can BQ25896 support I2C interface?", 
+#             "Does TDA2 vout support bt656 8-bit mode?"]
+
+# demo = gr.Interface(fn=qa_infer, inputs="text", allow_flagging='never', examples=EXAMPLES, cache_examples=False, outputs="text")
+# demo.launch()