Update app.py

app.py

@@ -1,49 +1,30 @@
 import streamlit as st
 from transformers import AutoTokenizer, TextGenerationPipeline
-from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
-import logging
+from auto_gptq import AutoGPTQForCausalLM
 
-# Load tokenizer and model
-device = "cuda" if st.checkbox("Use GPU") else "cpu"
+# Load the tokenizer and model
+pretrained_model_dir = "TheBloke/Llama-2-7b-Chat-GPTQ"
+quantized_model_dir = "amanchahar/llama2_finetune_Restaurants"
 
-model_name = "amanchahar/llama2_finetune_Restaurants"
-model = AutoGPTQForCausalLM.from_quantized(model_name, device="cuda:0", use_safetensors=True, use_triton=False)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(pretrained_model_dir, use_fast=True)
+model = AutoGPTQForCausalLM.from_quantized(quantized_model_dir, device="cuda:0")
 
-# Function to generate model response
-def generate_response(messages):
-    encodeds = tokenizer.apply_chat_template(messages, return_tensors="pt")
-    model_inputs = encodeds.to(device)
-    generated_ids = model.generate(model_inputs, max_new_tokens=1000, do_sample=True)
-    decoded_response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
-    return decoded_response[0]
+# Create a text generation pipeline
+pipeline = TextGenerationPipeline(model=model, tokenizer=tokenizer)
 
-# Streamlit app
+# Define the Streamlit app
 def main():
-    st.title("Interactive Conversational AI")
+    st.title("Restaurants Auto-GPTQ Text Generation")
 
-    messages = []
+    # User input text box
+    user_input = st.text_input("Enter your query:", "auto-gptq is")
 
-    while True:
-        # Display chat interface
-        user_input = st.text_input("You:", key="user_input")
-        if st.button("Send"):
-            if user_input.strip():
-                # Add user message to chat history
-                messages.append({"role": "user", "content": user_input})
-
-                # Generate response from the model
-                response = generate_response(messages)
-
-                # Display assistant's response
-                st.text_area("Assistant:", value=response, height=150)
-
-                # Add assistant's response to chat history
-                messages.append({"role": "assistant", "content": response})
+    if st.button("Generate"):
+        # Generate response based on user input
+        generated_text = pipeline(user_input)[0]["generated_text"]
+        st.markdown(f"**Generated Response:** {generated_text}")
 
 if __name__ == "__main__":
     main()
 
-
-