what-arm / app.py
gokul00060's picture
Create app.py
fe7c25c
raw
history blame
677 Bytes
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']}")