File size: 6,107 Bytes
a397aee
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
# import gradio as gr
# from huggingface_hub import InferenceClient

# """
# For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
# """
# client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")


# def respond(
#     message,
#     history: list[tuple[str, str]],
#     system_message,
#     max_tokens,
#     temperature,
#     top_p,
# ):
#     messages = [{"role": "system", "content": system_message}]

#     for val in history:
#         if val[0]:
#             messages.append({"role": "user", "content": val[0]})
#         if val[1]:
#             messages.append({"role": "assistant", "content": val[1]})

#     messages.append({"role": "user", "content": message})

#     response = ""

#     for message in client.chat_completion(
#         messages,
#         max_tokens=max_tokens,
#         stream=True,
#         temperature=temperature,
#         top_p=top_p,
#     ):
#         token = message.choices[0].delta.content

#         response += token
#         yield response

# """
# For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
# """
# demo = gr.ChatInterface(
#     respond,
#     additional_inputs=[
#         gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
#         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
#         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
#         gr.Slider(
#             minimum=0.1,
#             maximum=1.0,
#             value=0.95,
#             step=0.05,
#             label="Top-p (nucleus sampling)",
#         ),
#     ],
# )


# if __name__ == "__main__":
#     demo.launch()




import gradio as gr
import torch
from transformers import RagRetriever, RagSequenceForGeneration, AutoTokenizer

"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def strip_title(title):
    if title.startswith('"'):
        title = title[1:]
    if title.endswith('"'):
        title = title[:-1]
    return title

def retrieved_info(rag_model, query):
    # Tokenize query
    retriever_input_ids = rag_model.retriever.question_encoder_tokenizer.batch_encode_plus(
        [query],
        return_tensors="pt",
        padding=True,
        truncation=True,
    )["input_ids"].to(device)

    # Retrieve documents
    question_enc_outputs = rag_model.rag.question_encoder(retriever_input_ids)
    question_enc_pool_output = question_enc_outputs[0]

    result = rag_model.retriever(
        retriever_input_ids,
        question_enc_pool_output.cpu().detach().to(torch.float32).numpy(),
        prefix=rag_model.rag.generator.config.prefix,
        n_docs=rag_model.config.n_docs,
        return_tensors="pt",
    )

    # Display retrieved documents including URLs
    all_docs = rag_model.retriever.index.get_doc_dicts(result.doc_ids)
    retrieved_context = []
    for docs in all_docs:
        titles = [strip_title(title) for title in docs["title"]]
        texts = docs["text"]
        for title, text in zip(titles, texts):
            #print(f"Title: {title}")
            #print(f"Context: {text}")
            retrieved_context.append(f"{title}: {text}")

    answer = retrieved_context
    return answer




def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens ,
    temperature,
    top_p,
):
    # Load model
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    dataset_path = "./sample/my_knowledge_dataset"
    index_path = "./sample/my_knowledge_dataset_hnsw_index.faiss"

    tokenizer = AutoTokenizer.from_pretrained("facebook/rag-sequence-nq")
    retriever = RagRetriever.from_pretrained("facebook/rag-sequence-nq", index_name="custom",
                                             passages_path = dataset_path,
                                             index_path = index_path,
                                             n_docs = 5)
    rag_model = RagSequenceForGeneration.from_pretrained('facebook/rag-sequence-nq', retriever=retriever)
    rag_model.retriever.init_retrieval()
    rag_model.to(device)

    if message:  # If there's a user query
        response = retrieved_info(rag_model, message)  # Get the answer from your local FAISS and Q&A model
        return response[0]

    # In case no message, return an empty string
    return ""



"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
# Custom title and description
title = "🧠 Welcome to Your AI Knowledge Assistant"
description = """
HI!!, I am a chatbot, I retrieves relevant information from a custom dataset using RAG. Ask any question, and let me assist you.
My capabilities and knowledge is limited right now because of computational resources. Originally I can acess more than a million files 
from my knowledge-base but, right now, I am limited to less than 1000 files. LET'S BEGGINNNN......
"""

demo = gr.ChatInterface(
    respond,
    type = 'messages',
    additional_inputs=[
        gr.Textbox(value="You are a helpful and friendly assistant.", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ],
    title=title,
    description=description,
    textbox=gr.Textbox(placeholder=["'What is the future of AI?' or 'App Development'"]),
    examples=[["✨Future of AI"], ["📱App Development"]],
    example_icons=["🤖", "📱"],
    theme="compact",
)


if __name__ == "__main__":
    demo.launch(share = True )