import streamlit as st
from langchain_huggingface import ChatHuggingFace, HuggingFacePipeline

@st.cache_resource
def load_llm():
    llm = HuggingFacePipeline.from_model_id(
        model_id="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
        task="text-generation",
        pipeline_kwargs=dict(
            temperature=0.5,
            max_new_tokens=150
        )
    )
    return ChatHuggingFace(llm=llm)

model=load_llm()

st.header("QnA Tool")

user_input = st.text_input("Enter your prompt")

if st.button("Summarize"):
    result = model.invoke(user_input)
    st.write(result.content)