import streamlit as st
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load your dataset from Hugging Face
dataset = load_dataset("diylocals/TestData")  # Replace with your actual username and dataset name

# Load the IBM Granite model and tokenizer
model_name = "ibm-granite/granite-3.0-8b-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Streamlit app title
st.title("IBM Granite Model Analysis")

# Input text area for user input
user_input = st.text_area("Enter text for analysis (e.g., voltage readings):", "")

if st.button("Analyze"):
    if user_input:
        # Prepare input for the model
        inputs = tokenizer(user_input, return_tensors="pt")

        # Generate output using the model
        outputs = model.generate(**inputs)

        # Decode and display output
        output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        st.write("Model Output:")
        st.write(output_text)
    else:
        st.warning("Please enter some text for analysis.")