Update app.py

app.py

@@ -1,56 +1,24 @@
-# app.py
-import os
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-from langchain.llms import HuggingFacePipeline
-from langchain import PromptTemplate, LLMChain
+from model_utils import load_model, create_pipeline, generate_text
 
-# Load the model and tokenizer
-model_name = "Giang07/Llama-2-7b-chat-QLoRa"
-
-model_dir = "/home/user/model"
-
-# Check if model is already downloaded
-config_path = os.path.join(model_dir, "config.json")
-# if not os.path.exists(config_path):
-    st.write("Downloading model...")
-#     os.makedirs(model_dir, exist_ok=True)
-    model = AutoModelForCausalLM.from_pretrained(model_name)
-#     model.save_pretrained(model_dir)
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-#     tokenizer.save_pretrained(model_dir)
-# else:
-    # st.write("Loading model from cache...")
-    # model = AutoModelForCausalLM.from_pretrained(model_dir)
-    # tokenizer = AutoTokenizer.from_pretrained(model_dir)
-
-# Create the pipeline
-hf_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
-
-# Define the Hugging Face Pipeline LLM
-llm = HuggingFacePipeline(pipeline=hf_pipeline)
-
-# Define a prompt template
-prompt_template = PromptTemplate(
-    input_variables=["input_text"],
-    template="Translate the following English text to French: {input_text}"
-)
-
-# Create the LLM chain
-llm_chain = LLMChain(prompt_template=prompt_template, llm=llm)
+# Load model and tokenizer
+st.write("Loading model...")
+model, tokenizer = load_model()
+hf_pipeline = create_pipeline(model, tokenizer)
+st.write("Model loaded successfully!")
 
 # Streamlit application
-st.title("Ti6 Llama2-QLoRa Test")
+st.title("Hugging Face Model with LangChain and Streamlit")
 
 # Text input field
-input_text = st.text_input("What is your words:")
+input_text = st.text_input("Enter text to translate:")
 
 # Button to generate response
 if st.button("Translate"):
     if input_text:
         # Generate response using LangChain
-        response = llm_chain.run({"input_text": input_text})
+        response = generate_text(hf_pipeline, input_text)
         # Display the response
-        st.write("Answer:", response)
+        st.write("Translated text:", response)
     else:
         st.write("Please enter some text.")