model_on_cli.py

from langchain_community.document_loaders import PyPDFLoader, DirectoryLoader
from langchain.prompts import PromptTemplate
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_community.llms import CTransformers
from langchain.chains import RetrievalQA
import os

# Setup paths
DB_FAISS_PATH = r"db_faiss"

# Custom prompt template for QA
custom_prompt_template = """you are a medical chatbot made by Episafe and dont ever mention meta or llama ,Use the following pieces of information to answer the user's question, use all the informations and try to find patterns for the response
If you don't know the answer,don't try to make up an answer, make the answer very concise

Context: {context}
Question: {question}

IMPORTANT :if i ask you who are you or any kind of these questions that concerns you,say you are a Medical bot developped by Episafe.
Only return the helpful answer below and nothing else.
Helpful answer:
"""

def set_custom_prompt():
    """
    Prompt template for QA retrieval for each vectorstore
    """
    prompt = PromptTemplate(template=custom_prompt_template,
                            input_variables=['context', 'question'])
    return prompt

# Retrieval QA Chain
def retrieval_qa_chain(llm, prompt, db):
    qa_chain = RetrievalQA.from_chain_type(llm=llm,
                                           chain_type='stuff',
                                           retriever=db.as_retriever(
                                               search_kwargs={'k': 2}),
                                           return_source_documents=True,
                                           chain_type_kwargs={'prompt': prompt}
                                           )
    return qa_chain

# Loading the model
def load_llm():
    # Load the locally downloaded model here
    llm = CTransformers(
        model=r"llama-2-7b-chat.ggmlv3.q8_0.bin",
        model_type="llama",
        max_new_tokens=1024,
        temperature=0.1,
        context_length=1024,
        
    )
    return llm

import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(torch.cuda.is_available())
# QA Model Function
def qa_bot():
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2",
                                       model_kwargs={'device': device})
    db = FAISS.load_local(DB_FAISS_PATH, embeddings)
    llm = load_llm()
    qa_prompt = set_custom_prompt()
    qa = retrieval_qa_chain(llm, qa_prompt, db)
    return qa

# Initialize QA Bot outside of the main function to avoid reloading for every query
qa_instance = qa_bot()

def final_result(query):
    response = qa_instance({'query': query})
    # Process and print the response properly depending on its structure
    return response

def main():
    print("Welcome to the EpiSafe Medical Bot.\nType 'exit' to quit.")
    while True:
        user_input = input("\nYou  :")
        if user_input.lower() == 'exit':
            print("Exiting the chatbot. Goodbye!")
            break
        response = final_result(user_input)
        print("Response:", response['result'])

if __name__ == "__main__":
    main()