from datasets import load_dataset
import streamlit as st
from huggingface_hub import hf_hub_download
import gzip
import json

@st.cache(allow_output_mutation=True)
def get_model():
    model_id = "srihariEmids/emidsinfo-fine-tune-llamamodel"
    model = hf_hub_download(model_id, library_name="transformers")
    return model

def predict(text):
    model = get_model()
    inputs = model.prepare_inputs_for_generation(text, max_length=512, return_tensors="pt")
    outputs = model.generate(**inputs)
    return outputs

def main():
    text = st.text_input("Enter text to analyze:")
    prediction = predict(text)
    st.write(prediction)

if __name__ == "__main__":
    main()