Flash112/Document_Extractor

1import streamlit as st2import os3from langchain.llms import HuggingFaceHub4from langchain_groq import ChatGroq5from langchain.chains import LLMChain6from langchain.prompts import PromptTemplate7 8# Load API tokens from environment variables9HF_TOKEN = os.getenv("HF_TOKEN")10GROQ_API_KEY = os.getenv("GROQ_API_KEY")11 12# Streamlit UI13st.title("Occupation and Industry Extractor Chatbot")14st.write("Enter a remark, and the chatbot will extract the occupation and industry.")15 16# Model selection17model_choice = st.selectbox("Select Model:", ["Hugging Face", "Groq API"])18 19@st.cache_resource20def load_hf_model():21    return HuggingFaceHub(repo_id="microsoft/Phi-3.5-mini-instruct", huggingfacehub_api_token=HF_TOKEN)  # Smaller model22 23@st.cache_resource24def load_groq_model():25    return ChatGroq(model_name="qwen-2.5-32b", groq_api_key=GROQ_API_KEY)26 27llm = load_hf_model() if model_choice == "Hugging Face" else load_groq_model()28 29# Define prompt template30prompt_template = PromptTemplate(31    input_variables=["remark"],32    template="""33    Extract the primary occupation and its industry from the following remark:34    Remark: {remark}35    Provide the response in the following structured format:36    Primary Occupation: <Extracted Primary Occupation>37    Industry of Primary Occupation: <Extracted Industry>38    """39)40 41# Create LLM Chain42llm_chain = LLMChain(llm=llm, prompt=prompt_template)43 44remark = st.text_area("Enter your remark:")45 46if st.button("Extract Information"):47    if remark.strip():48        response = llm_chain.run(remark)49        50        # Extract and format response51        st.subheader("Extracted Information")52        response_lines = response.split("\n")53        occupation = "Not found"54        industry = "Not found"55        56        for line in response_lines:57            if line.startswith("Primary Occupation:"):58                occupation = line.replace("Primary Occupation:", "").strip()59            elif line.startswith("Industry of Primary Occupation:"):60                industry = line.replace("Industry of Primary Occupation:", "").strip()61        62        st.write(f"**Primary Occupation:** {occupation}")63        st.write(f"**Industry of Primary Occupation:** {industry}")64    else:65        st.warning("Please enter a valid remark.")66