|
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 |
|
|
|
|
|
|
|
HF_TOKEN = os.environ.get("HF_TOKEN", None) |
|
|
|
|
|
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 |
|
|
|
|
|
|
|
|
|
model_id="mistralai/Mistral-7B-Instruct-v0.2" |
|
device = f'cuda:{cuda.current_device()}' if cuda.is_available() else 'cpu' |
|
|
|
|
|
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) |
|
|
|
|
|
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()]) |
|
|
|
|
|
generate_text = pipeline( |
|
model=model, |
|
tokenizer=tokenizer, |
|
return_full_text=True, |
|
task='text-generation', |
|
|
|
temperature=0.1, |
|
max_new_tokens=2048, |
|
|
|
) |
|
|
|
llm = HuggingFacePipeline(pipeline=generate_text) |
|
|
|
|
|
try: |
|
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 |
|
|
|
|
|
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. Consider all source documents: |
|
- "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. |
|
|
|
Example 1: |
|
{{ |
|
"question": "How to use IPU1_0 instead of A15_0 to process NDK in TDA2x-EVM", |
|
"answer": "To use IPU1_0 instead of A15_0 to process NDK in TDA2x-EVM, you need to modify the configuration file of the NDK application. Specifically, change the processor reference from 'A15_0' to 'IPU1_0'.", |
|
"related_questions": [ |
|
{{ |
|
"question": "Can you provide MLBP documentation on TDA2?", |
|
"answer": "MLB is documented for DRA devices in the TRM book, chapter 24.12." |
|
}}, |
|
{{ |
|
"question": "Hi, could you share me the TDA2x documents about Security(SPRUHS7) and Cryptographic(SPRUHS8) addendums?", |
|
"answer": "Most of TDA2 documents are on ti.com under the product folder." |
|
}}, |
|
{{ |
|
"question": "Is any one can provide us a way to access CDDS for nessary docs?", |
|
"answer": "Which document are you looking for?" |
|
}}, |
|
{{ |
|
"question": "What can you tell me about the TDA2 and TDA3 processors? Can they / do they run Linux?", |
|
"answer": "We have moved your post to the appropriate forum." |
|
}} |
|
] |
|
}} |
|
|
|
Final Answer: To use IPU1_0 instead of A15_0 to process NDK in TDA2x-EVM, you need to modify the configuration file of the NDK application. Specifically, change the processor reference from 'A15_0' to 'IPU1_0'. |
|
|
|
Example 2: |
|
{{ |
|
"question": "Can BQ25896 support I2C interface?", |
|
"answer": "Yes, the BQ25896 charger supports the I2C interface for communication.", |
|
"related_questions": [ |
|
{{ |
|
"question": "What are the main features of BQ25896?", |
|
"answer": "The BQ25896 features include high-efficiency, fast charging capability, and a wide input voltage range." |
|
}}, |
|
{{ |
|
"question": "How to configure the BQ25896 for USB charging?", |
|
"answer": "To configure the BQ25896 for USB charging, set the input current limit and the charging current via I2C registers." |
|
}} |
|
] |
|
}} |
|
|
|
Final Answer: Yes, the BQ25896 charger supports the I2C interface for communication. |
|
|
|
""" |
|
|
|
|
|
return prompt |
|
|
|
|
|
def qa_infer(query): |
|
formatted_prompt = format_prompt(query) |
|
result = chain({"question": formatted_prompt, "chat_history": chat_history}) |
|
for doc in result['source_documents']: |
|
print("-"*50) |
|
print("Retrieved Document:", doc.page_content) |
|
print("#"*100) |
|
print(result['answer']) |
|
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() |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|