Update app.py

app.py

@@ -1,6 +1,5 @@
 import streamlit as st
-from langchain_ollama import OllamaLLM
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 
 def main():
     # Set up the page
@@ -20,19 +19,19 @@ def main():
     # Input for user-provided data
     prompt = st.text_area("Enter the prompt here:")
 
-    # Initialize the models
-    gemma_model = OllamaLLM(model='gemma:2b')
+    # Load model directly
     tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b-it")
     model = AutoModelForCausalLM.from_pretrained("google/gemma-2-2b-it")
 
+    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+
     # Button to generate the nudge
     if st.button("Generate Nudge"):
-        if S_boss.strip():
+        if prompt.strip():
             with st.spinner("Generating nudges..."):
-                # Generate the response using Ollama LLM
-                response = gemma_model.invoke(input=f"I want you to analyze the {prompt}. Which contains top 3 strengths or weaknesses of a person being assessed. You will generate nudges for improving upon these strengths or fixing upon these weaknesses. If you don't find any data, just respond as - No data available.")
+                response = pipe(messages)
                 st.success("Nudges generated successfully!")
-                st.text_area("Generated Nudges:", response, height=200)
+                st.text_area("Pipeline Response:", str(response), height=200)
         else:
             st.warning("Please enter data to generate nudges.")