Update app.py

app.py

@@ -1,10 +1,10 @@
+
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import pipeline
 
-# Load your text generation model and tokenizer
-model_name = "mostafaHaydar/model_test"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+# Load the pipeline for text generation
+model_name = "mostafaHaydar/model_test"  # Replace with your model's name
+generator = pipeline("text-generation", model=model_name)
 
 # Set up the Streamlit app
 st.title("Text Generation with LLaMA 3")
@@ -15,10 +15,8 @@ prompt = st.text_area("Enter your prompt:")
 # Generate text when the user clicks the button
 if st.button("Generate"):
     if prompt:
-        # Tokenize and generate text
-        inputs = tokenizer(prompt, return_tensors="pt")
-        output = model.generate(**inputs, max_length=150)  # Adjust max_length as needed
-        generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+        # Generate text using the pipeline
+        generated_text = generator(prompt, max_length=150, num_return_sequences=1)[0]['generated_text']
 
         # Display the generated text
         st.subheader("Generated Text:")