#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 = "Xenova/gpt-3.5-turbo"

# 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)

llm = ChatOpenAI(model_name="gpt-3.5-turbo", openai_api_key="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")
#llm = ChatOpenAI(model_name="gpt-3.5-turbo", openai_api_key="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")
#llm = ChatOpenAI(model_name="gpt-3.5-turbo", openai_api_key="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")

# Initialize global variable for conversation memory
buffer_memory = ConversationBufferWindowMemory(k=8)
#buffer_memory = ConversationBufferWindowMemory(k=3)

conversation = ConversationChain(
    llm=llm,
    memory=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.
You respond in a short and 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,
)

# Launch Gradio Interface
iface.launch()

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