from transformers import pipeline
import gradio as gr

pipe=pipeline('text-generation',model="Arjs/Llama-2-7b-chatbot-finetune")
def respond(system_message):
    if "irctc" not in system_message:
        system_message+=" in irctc"
    result = pipe(f"<s>[INST] {system_message} [/INST]")
    yield result[0]['generated_text']
"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""


demo = gr.ChatInterface(
    respond,
)


if __name__ == "__main__":
    demo.launch()