Spaces:
Sleeping
Sleeping
File size: 5,473 Bytes
ac1eff7 bc0db31 ac1eff7 f245e68 ac1eff7 f245e68 ac1eff7 f245e68 ac1eff7 f245e68 ac1eff7 |
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 |
import os
import gradio as gr
import cohere
from typing import Generator
from langchain_chroma import Chroma
# from langchain_huggingface import HuggingFaceEmbeddings
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from langchain.schema.document import Document
from typing import List
class HFSpaceChatBot:
"""
A chatbot powered by Retrieval Augmented Generation (RAG) aimed
to be deployed on the Hugging Face Space platform.
"""
def __init__(self,
embedding_model_path: str,
vector_database_path: str,
top_k: int = 10,
embedding_model_name: str = os.getenv("EMBEDDING_MODEL"),
api_key: str = os.getenv("CO_API_KEY"),
device: str = os.getenv("DEVICE"),
system_prompt: str = "Answer the user's question",
**kwargs) -> None:
"""
Constructor for the HFSpaceChatBot class.
Args:
embedding_model_path (str): The path to the embedding model.
vector_database_path (str): The path to the vector database.
top_k (int): The number of top documents to retrieve.
embedding_model_name (str): The name of the embedding model.
api_key (str): The API key for the cohere API.
device (str): The device to run the model on.
system_prompt (str): The system prompt for the chatbot.
**kwargs: Additional keyword arguments (for the cohere API)
"""
self.chat_history = []
self.cclient = cohere.Client(api_key=api_key)
self.embedding_model = HuggingFaceEmbeddings(
model_name=embedding_model_name,
model_kwargs={"device": device},
encode_kwargs={"normalize_embeddings": True},
cache_folder=embedding_model_path
)
self.vector_database = Chroma(
persist_directory=vector_database_path,
embedding_function=self.embedding_model
)
self.top_k = top_k
self.system_prompt = system_prompt
self.model_params = kwargs
def _get_relevant_information(self,
message: str) -> List[Document]:
"""
Get the relevant information from the chat history.
Args:
message (str): The message to search for.
Returns:
List[Document]: A list of relevant documents.
"""
return self.vector_database.similarity_search(message, self.top_k)
def _fetch_response(self,
message: str,
*args) -> Generator[str, None, None]:
"""
Fetch the reponse from the cohere API.
Args:
message (str): The message of the user.
Returns:
Generator[str, None, None]: A generator yielding the output tokens.
"""
docs = self._get_relevant_information(message)
relevant_information = "\n".join(
[doc.page_content
for doc in docs])
final_message = f"{self.system_prompt}\nWith the help of the\
following context:\n{relevant_information}\n\
Answer the following question:\n{message}"
response = self.cclient.chat_stream(
message=final_message,
chat_history=self.chat_history,
**self.model_params
)
current_text = ""
for event in response:
if event.event_type == "text-generation":
current_text += event.text
yield current_text
self.chat_history.append({
"role": "USER",
"text": message
})
self.chat_history.append({
"role": "CHATBOT",
"text": current_text
})
def launch(self,
title: str,
description: str) -> None:
"""
Launch the chat interface.
Args:
title (str): The title of the chat interface.
description (str): The description of the chat interface.
"""
gr.ChatInterface(
fn=self._fetch_response,
title=title,
description=description
).launch()
# if __name__ == "__main__":
embedding_model_path = os.path.join(os.getcwd(), "model")
system_prompt = """You are now assuming the role of the personal assistant
of Ilan ALIOUCHOUCHE, a French Computer Science student.
Your task is to assist users by answering their
questions about Ilan. You have access to comprehensive
details about Ilan's education, skills, professional
experience, and interests.
""".replace("\n", "")
chatbot = HFSpaceChatBot(
embedding_model_path=embedding_model_path,
vector_database_path=os.path.join(os.getcwd(), "chromadb"),
system_prompt=system_prompt,
temperature=0.4
)
title = "🤖 Ilan's Personal Agent 🤖"
description = """
You can ask my assistant (almost) anything about me! :D
You are currently using the Hugging Face Space version 🤗. A Docker image 🐳 for local use, utilizing a GGUF model is also available [here](https://github.com/ilanaliouchouche/my-ai-cv/pkgs/container/my-cv)
"""
chatbot.launch(
title=title,
description=description
)
|