import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Load the model and tokenizer
model_name = "bragour/Camel-7b-chat-awq"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Function to generate responses
def generate_response(user_input, chat_history=[]):
    new_user_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt')
    bot_input_ids = torch.cat([torch.LongTensor(chat_history), new_user_input_ids], dim=-1) if chat_history else new_user_input_ids
    
    chat_history = model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
    
    response = tokenizer.decode(chat_history[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
    return response, chat_history.tolist()

# Gradio interface
def chat(user_input, history=[]):
    response, history = generate_response(user_input, history)
    return response, history

iface = gr.Interface(
    fn=chat, 
    inputs=[gr.inputs.Textbox(lines=7, label="Input Text"), gr.inputs.State()],
    outputs=[gr.outputs.Textbox(label="Response"), gr.outputs.State()],
    title="ChatBot",
    description="A simple chatbot using a pre-trained Camel-7b-chat model."
)

iface.launch()