Update app.py

app.py

@@ -1,20 +1,30 @@
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer
+import transformers
+import torch
 
-# Load model and tokenizer
-model_name = "meta-llama/Meta-Llama-3-8B"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+HF_TOKEN=st.secrets(["HF_Token"])
+# Load the model and pipeline
+model_id = "meta-llama/Meta-Llama-3-8B"
 
-def predict(input_text):
-    inputs = tokenizer(input_text, return_tensors="pt")
-    outputs = model.generate(**inputs)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+# Set up the pipeline with the Hugging Face token
+pipeline = transformers.pipeline(
+    "text-generation",
+    model=model_id,
+    model_kwargs={"torch_dtype": torch.bfloat16, "use_auth_token": HF_TOKEN},  # Pass the token here
+    device_map="auto"
+)
 
-# Streamlit interface
+# Streamlit user interface
 st.title("LLM Model Inference")
-input_text = st.text_area("Enter your text here:")
+input_text = st.text_input("Enter your prompt:")
+
 if st.button("Generate"):
-    output_text = predict(input_text)
-    st.write(output_text)
+    if input_text:  # Check if the input is not empty
+        # Generate text using the pipeline
+        response = pipeline(input_text, max_length=150, num_return_sequences=1)  # Customize as needed
+        st.write("Generated Response:")
+        st.write(response[0]['generated_text'])  # Display the generated text
+    else:
+        st.error("Please enter a prompt to generate text.")
+