from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers import pipeline
from streamlit import title, input_text, button, text

# Define model and tokenizer
model_name = "gokul00060/loora-chat-arm"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Create pipeline for inference
chat_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)

# Streamlit app
title("Chat with Lora Adaptor")

text_input = input_text("Enter your message:")

if button("Send"):
  # Generate response
  response = chat_pipeline(text_input, max_length=1024)
  text(f"Lora: {response[0]['text']}")