Update app.py

app.py

@@ -1,25 +1,31 @@
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import streamlit as st
-model = AutoModelForCausalLM.from_pretrained("RWKV/rwkv-4-169m-pile")
-tokenizer = AutoTokenizer.from_pretrained("RWKV/rwkv-4-169m-pile")
+model_id = "RWKV/rwkv-raven-1b5"
+model = AutoModelForCausalLM.from_pretrained(model_id).to(0)
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+
 def main():
-    st.title("Dragon Text Generator")
-    st.write("Enter a prompt and generate dragon-inspired text!")
+    st.title("Raven Text Generator")
+    st.write("Ask a question about ravens and get a response!")
 
-    # Input prompt
-    prompt = st.text_area("Enter your prompt", value="", height=150)
+    # Input question
+    question = st.text_input("Ask a question")
 
-    if st.button("Generate Text"):
-        if prompt.strip() != "":
-            # Generate text based on the provided prompt
-            inputs = tokenizer(prompt, return_tensors="pt")
-            output = model.generate(inputs["input_ids"], max_new_tokens=20)
-            generated_text = tokenizer.decode(output[0].tolist())
-            st.markdown("## Generated Text")
+    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").to(0)
+            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 prompt.")
+            st.warning("Please enter a question.")
+
+if __name__ == "__main__":
+    main()
 
 if __name__ == "__main__":
     main()