app.py

from transformers import AutoModelForCausalLM, AutoTokenizer

from langchain.chat_models import ChatOpenAI
from langchain.chains import ConversationChain
from langchain.chains.conversation.memory import ConversationBufferWindowMemory
from langchain.prompts import PromptTemplate

import gradio as gr

REPO_ID = "mathslearn/llama-2-7b-chat-TCMKB"

# Load the model and tokenizer from Hugging Face's model hub
model = AutoModelForCausalLM.from_pretrained(REPO_ID)
tokenizer = AutoTokenizer.from_pretrained(REPO_ID)

llm = ChatOpenAI(model=model, tokenizer=tokenizer)

if 'buffer_memory' not in gr.session_state:
    gr.session_state.buffer_memory = ConversationBufferWindowMemory(k=8)
    #gr.session_state.buffer_memory = ConversationBufferWindowMemory(k=3)

conversation = ConversationChain(
    llm=llm,
    memory=gr.session_state.buffer_memory,
    verbose=True
)

context = """ 
You act as a chatbot to interact with users on their questions about traditional chinese medicine (TCM). \
Welcome the user in a friendly way and state this disclaimer: \
It is important to note that a chatbot is not a substitute for medical advice \
from a qualified healthcare professional. \
You respond in a short, very friendly style. \
For each text, mark NER tags. \
Tag categories: location, product. \
Text: I want to visit a clinic in [Ang Mo Kio](location). \
If a user says he/she wants to visit a traditional chinese medicine (TCM) clinic, \
ask for which location he/she wants to visit. \
After user replies with location, show in chat window the Google Map from this link \
https://www.google.com/maps/search/tcm+clinics+at+location+name \
For example, if a user wants to visit a clinic in Ang Mo Kio, \
show in chat window the Google Map from this link \
https://www.google.com/maps/search/tcm+clinics+at+ang+mo+kio \
Examoples of location names:
North:
Sembawang
Woodlands
Yishun
North-East:
Ang Mo Kio
Hougang
Punggol
Sengkang
Serangoon
East:
Bedok
Pasir Ris
Tampines
West:
Bukit Batok
Bukit Panjang
Choa Chu Kang
Clementi
Jurong East
Jurong West
Tengah
Central:
Bishan
Bukit Merah
Bukit Timah
Central Area
Geylang
Kallang
Whampoa
Marine Parade
Queenstown
Toa Payoh
For each text, mark NER tags. \
Tag categories: location, product. \
Text: I want to buy/get [Po Chai Pills](product). \
If a user wants to buy/get a product, suggest that \
he/she can consider buying/getting from https://www.amazon.sg/s?k=product+name \
For example, if a user wants to buy Po Chai Pills, suggest \
he/she can consider buying/getting from https://www.amazon.sg/s?k=po+chai+pills \
Examples of product names:
Ointment/Hong You/Feng You/Fengyou
Liquorice/Gan cao/Gancao
Chrysanthemum/Ju hua/Juhua
Goji berry/wolfberry/Gou Qi Zi/Gouqizi
Red dates/Jujubes/Hong Zao/Hongzao
"""

prompt_template = PromptTemplate.from_template('''system role :{context} \
        user:{query}\
        assistance: 
        ''')

# Define Gradio Interface
iface = gr.Interface(
    fn=lambda query: conversation.run(prompt_template.format(context=context, query=query)),
    inputs=gr.Textbox(),
    outputs=gr.Textbox(),
    live=True,
    capture_session=True
)

# Launch Gradio Interface
iface.launch()

# gr.load("models/ksh-nyp/llama-2-7b-chat-TCMKB").launch()