File size: 2,690 Bytes
0394b1d
6dacbc2
 
 
 
 
0394b1d
6dacbc2
 
0394b1d
6dacbc2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0cf95f2
ceaa0db
 
0cf95f2
ceaa0db
52548bf
0cf95f2
52548bf
0cf95f2
6dacbc2
ceaa0db
6dacbc2
0cf95f2
6fea818
0cf95f2
6dacbc2
 
 
 
 
0cf95f2
6dacbc2
0cf95f2
6fea818
0cf95f2
6dacbc2
 
 
 
 
a798180
6dacbc2
0cf95f2
6dacbc2
 
 
 
 
 
 
 
 
 
0cf95f2
6dacbc2
c98d37d
0cf95f2
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
79
80
81
82
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_message(message: str, mode: str):
    
    # Check if query is empty
    if not message.strip():
        return "Enter a valid message."
    if mode == "Chat":
        return chat_mode(message)
    elif mode == "Web-Search":
        return web_search(message)
    else:
        return "Select a valid mode."

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

def web_search(message: str):
    global messages
    similar_docs = db.similarity_search(message, 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 give guidance relevant to the query.

    Query: {message}

    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: {message}\n\nIT-Assistant: {response.content}"

demo = gr.Interface(
    fn=handle_message,
    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()