from langchain.chains import ConversationChain
from transformers import AutoTokenizer, AutoModelForCausalLM
import gradio as gr
from myLLM import AutoModelLanguageModel

model_name = "microsoft/DialoGPT-medium"
tokenizer = AutoTokenizer.from_pretrained(model_name)
llm = AutoModelLanguageModel(model_name)

agent = ConversationChain(llm=llm)


# Define the Gradio interface
def chatbot_interface(input_text):
    response = agent(input_text)
    return response

# Define the Gradio app
gradio_app = gr.Interface(
    fn=chatbot_interface,
    inputs=gr.inputs.Textbox(prompt="Say something..."),
    outputs=gr.outputs.Textbox(),
    title="ConversationChain Chatbot",
    description="A chatbot interface powered by ConversationChain and Hugging Face.",
)

# Run the Gradio app
if __name__ == "__main__":
    gradio_app.run()