Update app.py

app.py

@@ -1,60 +1,60 @@
 import os
 import torch
 from torch import cuda, bfloat16
-from transformers import AutoTokenizer, pipeline, BitsAndBytesConfig
+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
-from huggingface_hub import InferenceClient
+
 
 # Load the Hugging Face token from environment
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 
-# Load the Mistral model and tokenizer
-model_id = 'mistralai/Mistral-7B-Instruct-v0.3'
-client = InferenceClient(model_id)
-
-# # Define stopping criteria
-# class StopOnTokens:
-#     def __call__(self, input_ids, scores, **kwargs):
-#         for stop_ids in stop_token_ids:
-#             if torch.eq(input_ids[0][-len(stop_ids):], stop_ids).all():
-#                 return True
-#         return False
-
-# # Define stopping criteria list
-# stop_list = ['\nHuman:', '\n```\n']
-# stop_token_ids = [client.tokenizer(x)['input_ids'] for x in stop_list]
-# stop_token_ids = [torch.LongTensor(x).to(cuda.current_device() if cuda.is_available() else 'cpu') for x in stop_token_ids]
+# 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'
+model_id= "meta-llama/Llama-2-7b-chat-hf"
+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
-def generate(prompt, history = None, system_prompt=None, temperature=0.2, max_new_tokens=1024, top_p=0.95, repetition_penalty=1.0):
-    temperature = float(temperature)
-    if temperature < 1e-2:
-        temperature = 1e-2
-    top_p = float(top_p)
-
-    generate_kwargs = dict(
-        temperature=temperature,
-        max_new_tokens=max_new_tokens,
-        top_p=top_p,
-        repetition_penalty=repetition_penalty,
-        do_sample=True,
-        seed=42,
-    )
-
-    formatted_prompt = format_prompt(prompt, history, system_prompt)
-    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
-    output = ""
-
-    for response in stream:
-        output += response.token.text
-        yield output
-    return output
-
-llm = HuggingFacePipeline(pipeline=generate)
+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:
@@ -103,3 +103,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()