from langchain.llms import CTransformers
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
import os
import io
import gradio as gr
import time

custom_prompt_template = """
You are an AI assistant and your task is to chat and provide information to the users
Below is the user's query.
Query: {query}
You return the helpful code and related details
Helpful code and related details:
Below are some examples of the inputs(instruction) and the outputs(response)
[
  { "inputs": "### Instruction: Introduce yourself. \n\n### Response: Hi, I'm a chatbot designed to assist and engage in conversations on various topics. How can I help you today?" },
  { "inputs": "### Instruction: What is your favourite hobby? \n\n### Response: As a chatbot, I don't have hobbies, but I'm here to chat about any topic you're interested in. What are your interests?" },
  { "inputs": "### Instruction: Tell me a joke. \n\n### Response: Why don't scientists trust atoms? Because they make up everything!" },
  { "inputs": "### Instruction: What is the capital of France? \n\n### Response: The capital of France is Paris. Anything else you'd like to know about?" },
  { "inputs": "### Instruction: Discuss the weather. \n\n### Response: The weather can vary greatly depending on your location. Can you specify a city or region you're interested in?" },
  { "inputs": "### Instruction: Explain the concept of artificial intelligence. \n\n### Response: Artificial Intelligence, or AI, refers to the development of computer systems that can perform tasks that typically require human intelligence. It includes areas like machine learning and natural language processing." },
  { "inputs": "### Instruction: What are your thoughts on technology? \n\n### Response: I don't have personal thoughts, but I can provide information and answer questions about technology. What aspect of technology interests you?" },
  { "inputs": "### Instruction: How can I stay motivated? \n\n### Response: Staying motivated can be challenging. Setting clear goals, breaking them into smaller tasks, and celebrating achievements along the way can help. What are you working on?" },
  { "inputs": "### Instruction: Share a fun fact. \n\n### Response: Did you know that honey never spoils? Archaeologists have found pots of honey in ancient Egyptian tombs that are over 3,000 years old and still perfectly edible!" },
  { "inputs": "### Instruction: What's the best way to learn a new language? \n\n### Response: Learning a new language involves consistent practice, immersion, and using a variety of resources like books, apps, and conversation partners. What language are you interested in learning?" }
]

"""

def set_custom_prompt():
    prompt = PromptTemplate(
        template = custom_prompt_template,
        input_variables = ['query']
    )
    return prompt

def load_model():
    llm = CTransformers(
        model = 'TheBloke/Llama-2-7B-GGUF',
        model_type = 'llama',
        max_new_tokens = 1096,
        temperature = 0.6,
        repetition_penalty = 1.13,
        gpu_layers = 2
    )
    return llm

def chain_pipeline():
    llm = load_model()
    qa_prompt = set_custom_prompt()
    qa_chain = LLMChain(
        prompt = qa_prompt,
        llm=llm
    )
    return qa_chain

llmcahin = chain_pipeline()

def bot(query):
    llm_response = llmcahin.run({"query":query})
    return llm_response

with gr.Blocks(title="chat llama 7b") as demo:
    gr.Markdown("# chat llama")
    chatbot = gr.Chatbot([],elem_id="chatbot",height=700)
    msg = gr.Textbox()
    clear = gr.ClearButton([msg,chatbot])

    def respond(message, chat_history):
        bot_message = bot(message)
        chat_history.append((message, bot_message))
        time.sleep(2)
        return "",chat_history
    
    msg.submit(respond,[msg, chatbot],[msg, chatbot])

demo.launch(share=True)