|
from dotenv import load_dotenv |
|
import gradio as gr |
|
import os |
|
from llama_index.core import StorageContext, load_index_from_storage, VectorStoreIndex, SimpleDirectoryReader, ChatPromptTemplate, Settings |
|
from llama_index.llms.huggingface import HuggingFaceInferenceAPI |
|
from llama_index.embeddings.huggingface import HuggingFaceEmbedding |
|
from sentence_transformers import SentenceTransformer |
|
import firebase_admin |
|
from firebase_admin import db, credentials |
|
import datetime |
|
import uuid |
|
import random |
|
|
|
def select_random_name(): |
|
names = ['Clara', 'Lily'] |
|
return random.choice(names) |
|
|
|
|
|
|
|
load_dotenv() |
|
|
|
cred = credentials.Certificate("redfernstech-fd8fe-firebase-adminsdk-g9vcn-0537b4efd6.json") |
|
firebase_admin.initialize_app(cred, {"databaseURL": "https://redfernstech-fd8fe-default-rtdb.firebaseio.com/"}) |
|
|
|
Settings.llm = HuggingFaceInferenceAPI( |
|
model_name="meta-llama/Meta-Llama-3-8B-Instruct", |
|
tokenizer_name="meta-llama/Meta-Llama-3-8B-Instruct", |
|
context_window=3000, |
|
token=os.getenv("HF_TOKEN"), |
|
max_new_tokens=512, |
|
generate_kwargs={"temperature": 0.1}, |
|
) |
|
Settings.embed_model = HuggingFaceEmbedding( |
|
model_name="BAAI/bge-small-en-v1.5" |
|
) |
|
|
|
|
|
PERSIST_DIR = "db" |
|
PDF_DIRECTORY = 'data' |
|
|
|
|
|
os.makedirs(PDF_DIRECTORY, exist_ok=True) |
|
os.makedirs(PERSIST_DIR, exist_ok=True) |
|
|
|
|
|
current_chat_history = [] |
|
|
|
def data_ingestion_from_directory(): |
|
|
|
documents = SimpleDirectoryReader(PDF_DIRECTORY).load_data() |
|
storage_context = StorageContext.from_defaults() |
|
index = VectorStoreIndex.from_documents(documents) |
|
index.storage_context.persist(persist_dir=PERSIST_DIR) |
|
|
|
def handle_query(query): |
|
chat_text_qa_msgs = [ |
|
( |
|
"user", |
|
""" |
|
You're Lily, working as customer care at RedfernsTech. Continue the conversation flow and give responses within 10-15 words only. |
|
{context_str} |
|
Question: |
|
{query_str} |
|
""" |
|
) |
|
] |
|
text_qa_template = ChatPromptTemplate.from_messages(chat_text_qa_msgs) |
|
|
|
|
|
storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR) |
|
index = load_index_from_storage(storage_context) |
|
|
|
|
|
context_str = "" |
|
for past_query, response in reversed(current_chat_history): |
|
if past_query.strip(): |
|
context_str += f"User asked: '{past_query}'\nBot answered: '{response}'\n" |
|
|
|
query_engine = index.as_query_engine(text_qa_template=text_qa_template, context_str=context_str) |
|
answer = query_engine.query(query) |
|
|
|
if hasattr(answer, 'response'): |
|
response = answer.response |
|
elif isinstance(answer, dict) and 'response' in answer: |
|
response = answer['response'] |
|
else: |
|
response = "Sorry, I couldn't find an answer." |
|
|
|
|
|
current_chat_history.append((query, response)) |
|
|
|
return response |
|
|
|
|
|
print("Processing PDF ingestion from directory:", PDF_DIRECTORY) |
|
data_ingestion_from_directory() |
|
|
|
|
|
"""def predict(message,history): |
|
response = handle_query(message) |
|
return response""" |
|
def predict(message, history): |
|
logo_html = ''' |
|
<div class="circle-logo"> |
|
<img src="https://rb.gy/8r06eg" alt="FernAi"> |
|
</div> |
|
''' |
|
response = handle_query(message) |
|
response_with_logo = f'<div class="response-with-logo">{logo_html}<div class="response-text">{response}</div></div>' |
|
return response_with_logo |
|
def save_chat_message(session_id, message_data): |
|
ref = db.reference(f'/chat_history/{session_id}') |
|
ref.push().set(message_data) |
|
|
|
|
|
def chat_interface(message, history): |
|
try: |
|
|
|
session_id = str(uuid.uuid4()) |
|
|
|
|
|
response = handle_query(message) |
|
|
|
|
|
message_data = { |
|
"sender": "user", |
|
"message": message, |
|
"response": response, |
|
"timestamp": datetime.datetime.now().isoformat() |
|
} |
|
|
|
|
|
save_chat_message(session_id, message_data) |
|
|
|
|
|
return response |
|
except Exception as e: |
|
return str(e) |
|
|
|
|
|
css = ''' |
|
.circle-logo { |
|
display: inline-block; |
|
width: 40px; |
|
height: 40px; |
|
border-radius: 50%; |
|
overflow: hidden; |
|
margin-right: 10px; |
|
vertical-align: middle; |
|
} |
|
.circle-logo img { |
|
width: 100%; |
|
height: 100%; |
|
object-fit: cover; |
|
} |
|
.response-with-logo { |
|
display: flex; |
|
align-items: center; |
|
margin-bottom: 10px; |
|
} |
|
footer { |
|
display: none !important; |
|
background-color: #F8D7DA; |
|
} |
|
label.svelte-1b6s6s {display: none} |
|
''' |
|
gr.ChatInterface(chat_interface, |
|
css=css, |
|
description="Clara", |
|
clear_btn=None, undo_btn=None, retry_btn=None, |
|
).launch() |