Update app.py

app.py

@@ -1,34 +1,22 @@
-from transformers import pipeline
 import streamlit as st
-import ollama
-from typing import Dict, Generator
+from ollama import Ollama
 
-def ollama_generator(model_name: str, messages: Dict) -> Generator:
-    stream = ollama.chat(
-        model=model_name, messages=messages, stream=True)
-    for chunk in stream:
-        yield chunk['message']['content']
+# Initialize Ollama model
+ollama = Ollama(model='gpt-3')
 
+st.title("Text Generation with Ollama and Streamlit")
 
-st.title("Ollama with Streamlit")
-if "selected_model" not in st.session_state:
-    st.session_state.selected_model = ""
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-st.session_state.selected_model = st.selectbox(
-    "Please select the model:", [model["name"] for model in ollama.list()["models"]])
-for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
-if prompt := st.chat_input("How could I help you?"):
-    # Add user message to chat history
-    st.session_state.messages.append({"role": "user", "content": prompt})
-    # Display user message in chat message container
-    with st.chat_message("user"):
-        st.markdown(prompt)
+# Input text from user
+prompt = st.text_area("Enter your prompt:", "")
 
-    with st.chat_message("assistant"):
-        response = st.write_stream(ollama_generator(
-            st.session_state.selected_model, st.session_state.messages))
-    st.session_state.messages.append(
-        {"role": "assistant", "content": response})
+if st.button("Generate Text"):
+    if prompt:
+        # Generate text using Ollama
+        response = ollama.generate(prompt)
+        st.text_area("Generated Text:", response, height=200)
+    else:
+        st.warning("Please enter a prompt.")
+
+
+if __name__ == "__main__":
+    main()