from transformers import AutoTokenizer, AutoModelForCausalLM
import transformers
import torch

# model = "tiiuae/falcon-40b-instruct"
model = "tiiuae/falcon-7b-instruct"

tokenizer = AutoTokenizer.from_pretrained(model, offload_folder="offload", offload_state_dict = True, )
pipeline = transformers.pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    # torch_dtype=torch.bfloat16,
    trust_remote_code=True,
    device_map="auto",
)

def format_chat_prompt(message, chat_history):
    prompt = ""
    for turn in chat_history:
        user_message, bot_message = turn
        prompt = f"{prompt}\nUser: {user_message}\nAssistant: {bot_message}"
    prompt = f"{prompt}\nUser: {message}\nAssistant:"
    return prompt

def respond(message, chat_history):
        formatted_prompt = format_chat_prompt(message, chat_history)
        # print(formatted_prompt)
        bot_message = generate_seqs(prompt = formatted_prompt,
                                     max_new_tokens=1024,
                                     stop_sequence=["\nUser:", "<|endoftext|>"]).split('Assistant: ')[-1]

        chat_history.append((message, bot_message))
        return "", chat_history

def generate_seqs(prompt, max_new_tokens=None, stop_sequence=None):
    output = pipeline(prompt,
      max_length=200,
      max_new_tokens = max_new_tokens,
      stop_sequence = stop_sequence,
      do_sample=True,
      top_k=10,
      num_return_sequences=1,
      eos_token_id=tokenizer.eos_token_id)
    return output[0]['generated_text']

with gr.Blocks() as demo:
    chatbot = gr.Chatbot(height=240) #just to fit the notebook
    msg = gr.Textbox(label="Prompt")
    btn = gr.Button("Submit")
    clear = gr.ClearButton(components=[msg, chatbot], value="Clear console")

    btn.click(respond, inputs=[msg, chatbot], outputs=[msg, chatbot])
    msg.submit(respond, inputs=[msg, chatbot], outputs=[msg, chatbot]) #Press enter to submit
demo.launch()