File size: 2,542 Bytes
0394b1d
6dacbc2
 
 
 
 
0394b1d
6dacbc2
 
0394b1d
6dacbc2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6ab4abb
52548bf
6dacbc2
52548bf
6dacbc2
 
 
 
 
6fea818
6dacbc2
 
 
 
 
 
7c44a86
6dacbc2
 
6fea818
6dacbc2
 
 
 
 
 
bb2fbe5
6dacbc2
 
 
 
 
 
 
 
 
 
 
 
7c44a86
6dacbc2
c98d37d
 
c2f639d
c98d37d
c2f639d
0394b1d
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
import gradio as gr
from langchain.embeddings import SentenceTransformerEmbeddings
from langchain.vectorstores import FAISS
from langchain_community.chat_models.huggingface import ChatHuggingFace
from langchain.schema import SystemMessage, HumanMessage, AIMessage
from langchain_community.llms import HuggingFaceEndpoint

model_name = "sentence-transformers/all-mpnet-base-v2"
embedding_llm = SentenceTransformerEmbeddings(model_name=model_name)

db = FAISS.load_local("faiss_index", embedding_llm, allow_dangerous_deserialization=True)

# Set up Hugging Face model
llm = HuggingFaceEndpoint(
    repo_id="HuggingFaceH4/starchat2-15b-v0.1",
    task="text-generation",
    max_new_tokens=4096,
    temperature=0.6,
    top_p=0.9,
    top_k=40,
    repetition_penalty=1.2,
    do_sample=True,
)
chat_model = ChatHuggingFace(llm=llm)

messages = [
    SystemMessage(content="You are a helpful assistant."),
    HumanMessage(content="Hi AI, how are you today?"),
    AIMessage(content="I'm great thank you. How can I help you?")
]

def handle_query(query: str, mode: str):
    if mode == "Chat":
        return chat_mode(query)
    elif mode == "Web-Search":
        return web_search(query)
    else:
        return "Invalid mode selected."

def chat_mode(query: str):
    global messages
    prompt = HumanMessage(content=query)
    messages.append(prompt)
    response = chat_model.invoke(messages)
    messages.append(response.content)
    if len(messages) >= 6:
        messages = messages[-6:]
    return f"You: {query}\nIT-Assistant: {response.content}"

def web_search(query: str):
    global messages
    similar_docs = db.similarity_search(query, k=3)
    if similar_docs:
        source_knowledge = "\n".join([x.page_content for x in similar_docs])
    else:
        source_knowledge = ""
    augmented_prompt = f"""
    If the answer to the next query is not contained in the Web Search say 'No Answer Available' and then give advice.

    Query: {query}

    Web Search:
    {source_knowledge}
    """
    prompt = HumanMessage(content=augmented_prompt)
    messages.append(prompt)
    response = chat_model.invoke(messages)
    messages.append(response.content)
    if len(messages) >= 6:
        messages = messages[-6:]
    return f"You: {query}\nIT-Assistant: {response.content}"

demo = gr.Interface(
    fn=handle_query,
    inputs=["text", gr.Radio(["Chat", "Web-Search"], label="Mode", info="Choose a mode and enter your message, then click submit to interact.")],
    outputs="text",
    title="IT Assistant")
demo.launch()