diff --git a/setup.py b/setup.py index 22066b3..74259b6 100755 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setuptools.setup( name="text2text", - version="1.7.4", + version="1.7.5", author="artitw", author_email="artitw@gmail.com", description="Text2Text: Crosslingual NLP/G toolkit", @@ -18,7 +18,7 @@ "License :: OSI Approved :: MIT License", "Operating System :: OS Independent", ], - keywords='multilingual crosslingual gpt chatgpt bert natural language processing nlp nlg text generation gpt question answer answering information retrieval tfidf tf-idf bm25 search index summary summarizer summarization tokenizer tokenization translation backtranslation data augmentation science machine learning colab embedding levenshtein sub-word edit distance conversational dialog chatbot llama', + keywords='multilingual crosslingual gpt chatgpt bert natural language processing nlp nlg text generation gpt question answer answering information retrieval tfidf tf-idf bm25 search index summary summarizer summarization tokenizer tokenization translation backtranslation data augmentation science machine learning colab embedding levenshtein sub-word edit distance conversational dialog chatbot llama rag', install_requires=[ 'faiss-cpu', 'flask', diff --git a/text2text/rag_assistant.py b/text2text/rag_assistant.py index 0f17ade..0024e2b 100644 --- a/text2text/rag_assistant.py +++ b/text2text/rag_assistant.py @@ -26,6 +26,21 @@ def is_valid_url(url): except Exception: return False +def is_affirmative(response): + affirmative_keywords = [ + "yes", "yeah", "yep", "sure", "absolutely", "definitely", + "certainly", "of course", "indeed", "affirmative", "correct", + "right", "exactly", "true", "positive" + ] + + response_lower = response.lower() + + for keyword in affirmative_keywords: + if keyword in response_lower: + return True + + return False + class RagAssistant(t2t.Assistant): def __init__(self, **kwargs): super().__init__(**kwargs) @@ -57,7 +72,17 @@ def __init__(self, **kwargs): def chat_completion(self, messages=[{"role": "user", "content": "hello"}], stream=False, schema=None, **kwargs): k = kwargs.get("k", 3) query = messages[-1]["content"] - docs = self.index.retrieve([query], k=k)[0] - grounding_information = "Ground your response based on the following information:\n\n" + "\n- ".join(docs) - messages[-1] = {"role": "user", "content": query + "\n\n" + grounding_information} + question_check = f"Respond YES if this is a question; otherwise respond NO: {query}" + question_check = [{"role": "user", "content": question_check}] + response = t2t.Assistant.chat_completion(self, question_check)["message"]["content"] + docs = [] + if is_affirmative(response): + reword_prompt = f"Reword this question to be a demand: {query}" + reword_prompt = [{"role": "user", "content": reword_prompt}] + demand = t2t.Assistant.chat_completion(self, reword_prompt)["message"]["content"] + docs = self.index.retrieve([demand], k=k)[0] + else: + docs = self.index.retrieve([query], k=k)[0] + grounding_prompt = "Base your response on the following information:\n\n" + "\n- ".join(docs) + messages[-1] = {"role": "user", "content": query + "\n\n" + grounding_prompt} return t2t.Assistant.chat_completion(self, messages=messages, stream=stream, schema=schema, **kwargs)