|
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 |
|
|
|
|
|
load_dotenv() |
|
|
|
|
|
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", |
|
""" |
|
As FernAI, your goal is to offer top-tier service and information about RedFerns Tech company. |
|
Provide answers naturally based on the conversation flow. Ultimately, aim to attract users to connect with our company services. . |
|
Summarize answers effectively in fewer words without unnecessary repetition. |
|
{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 |
|
|
|
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; |
|
} |
|
''' |
|
gr.ChatInterface(predict, |
|
css=css, |
|
description="FernAI", |
|
clear_btn=None, undo_btn=None, retry_btn=None, |
|
examples=['Tell me about Redfernstech?', 'Services in Redfernstech?'] |
|
).launch() |