import gradio as gr
from transformers import pipeline
from transformers import AutoModelForCausalLM

# chat = pipeline("text-generation", model="LeoLM/leo-mistral-hessianai-7b-chat")

model = AutoModelForCausalLM.from_pretrained("LeoLM/leo-mistral-hessianai-7b-chat", load_in_4bit=True, device_map="auto")


def chat_with_model(prompt, history):
    #response = chat(prompt, max_length=50)
    #return response[0]['generated_text']
    return "works"
demo = gr.ChatInterface(fn=chat_with_model, title="Echo Bot")
demo.launch()