RWKV / app.py
varun500's picture
Update app.py
a042655
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("samwit/koala-7b")
model = AutoModelForCausalLM.from_pretrained("samwit/koala-7b")
st.title("Raven Text Generator")
st.write("Ask a question about ravens and get a response!")
# Input question
question = st.text_input("Ask a question")
if st.button("Generate Response"):
if question.strip() != "":
# Generate response based on the provided question
prompt = f"### Instruction: {question}\n### Response:"
inputs = tokenizer(prompt, return_tensors="pt")
output = model.generate(inputs["input_ids"], max_new_tokens=100)
generated_text = tokenizer.decode(output[0].tolist(), skip_special_tokens=True)
st.markdown("## Generated Response")
st.write(generated_text)
else:
st.warning("Please enter a question.")