Update app.py

app.py

@@ -3,12 +3,12 @@ import os
 import uuid
 import threading
 import pandas as pd
-import torch
 from langchain.document_loaders.csv_loader import CSVLoader
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
-from langchain.llms import CTransformers
+from langchain.llms import HuggingFacePipeline
 from langchain.chains import ConversationalRetrievalChain
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
 # Global model cache
 MODEL_CACHE = {
@@ -20,20 +20,36 @@ MODEL_CACHE = {
 os.makedirs("user_data", exist_ok=True)
 
 def initialize_model_once():
-    """Initialize the model once and cache it"""
+    """Initialize model once using pipeline API"""
     with MODEL_CACHE["init_lock"]:
         if MODEL_CACHE["model"] is None:
-            # Path ke model local dalam repository
-            model_path = "tinyllama-1.1b-chat-v1.0.Q5_K_M.gguf"
-            MODEL_CACHE["model"] = CTransformers(
-                model=model_path,
-                model_type="tinyllama",
+            # Load model from Hugging Face Hub
+            model_id = "meta-llama/Llama-2-7b-chat-hf"
+            
+            # Tokenizer
+            tokenizer = AutoTokenizer.from_pretrained(model_id, token=os.environ.get("HF_TOKEN"))
+            
+            # Model with low precision
+            model = AutoModelForCausalLM.from_pretrained(
+                model_id, 
+                token=os.environ.get("HF_TOKEN"),
+                device_map="auto",
+                load_in_8bit=True  # Quantize model to 8-bit precision
+            )
+            
+            # Create pipeline
+            pipe = pipeline(
+                "text-generation",
+                model=model,
+                tokenizer=tokenizer,
                 max_new_tokens=512,
                 temperature=0.2,
                 top_p=0.9,
-                top_k=50,
                 repetition_penalty=1.2
             )
+            
+            # Create LangChain wrapper
+            MODEL_CACHE["model"] = HuggingFacePipeline(pipeline=pipe)
     
     return MODEL_CACHE["model"]
 
@@ -53,27 +69,18 @@ class ChatBot:
             # Handle file from Gradio
             file_path = file.name if hasattr(file, 'name') else str(file)
             
-            # Copy to user directory
-            user_file_path = f"{self.user_dir}/uploaded.csv"
-            
-            # For debugging
-            print(f"Processing file: {file_path}")
-            print(f"Saving to: {user_file_path}")
-            
-            # Verify the CSV can be loaded
+            # Verify and save CSV
             try:
                 df = pd.read_csv(file_path)
-                print(f"CSV verified: {df.shape[0]} rows, {len(df.columns)} columns")
-                
-                # Save a copy in user directory
+                user_file_path = f"{self.user_dir}/uploaded.csv"
                 df.to_csv(user_file_path, index=False)
+                print(f"CSV verified: {df.shape[0]} rows, {len(df.columns)} columns")
             except Exception as e:
                 return f"Error membaca CSV: {str(e)}"
             
             # Load document
             try:
-                loader = CSVLoader(file_path=file_path, encoding="utf-8", csv_args={
-                    'delimiter': ','})
+                loader = CSVLoader(file_path=file_path, encoding="utf-8", csv_args={'delimiter': ','})
                 data = loader.load()
                 print(f"Documents loaded: {len(data)}")
             except Exception as e:
@@ -84,7 +91,7 @@ class ChatBot:
                 db_path = f"{self.user_dir}/db_faiss"
                 embeddings = HuggingFaceEmbeddings(
                     model_name='sentence-transformers/all-MiniLM-L6-v2',
-                    model_kwargs={'device': 'cuda' if torch.cuda.is_available() else 'cpu'}
+                    model_kwargs={'device': 'auto'}
                 )
                 
                 db = FAISS.from_documents(data, embeddings)
@@ -104,11 +111,11 @@ class ChatBot:
             except Exception as e:
                 return f"Error creating chain: {str(e)}"
             
-            # Add basic file info to chat history for context
+            # Add file info to chat history
             file_info = f"CSV berhasil dimuat dengan {df.shape[0]} baris dan {len(df.columns)} kolom. Kolom: {', '.join(df.columns.tolist())}"
             self.chat_history.append(("System", file_info))
             
-            return "File CSV berhasil diproses! Anda dapat mulai chat dengan model Llama2."
+            return "File CSV berhasil diproses! Anda dapat mulai chat dengan model Llama 2."
         except Exception as e:
             import traceback
             print(traceback.format_exc())
@@ -119,29 +126,23 @@ class ChatBot:
             return "Mohon upload file CSV terlebih dahulu."
         
         try:
-            # Process the question with the chain
+            # Process with the chain
             result = self.chain({"question": message, "chat_history": self.chat_history})
             
-            # Update internal chat history
+            # Update chat history
             answer = result["answer"]
             self.chat_history.append((message, answer))
             
-            # Return just the answer for Gradio
             return answer
         except Exception as e:
             import traceback
             print(traceback.format_exc())
             return f"Error: {str(e)}"
 
-    def cleanup(self):
-        """Release resources when session ends"""
-        self.chain = None
-
+# UI Code dan handler functions sama seperti sebelumnya
 def create_gradio_interface():
     with gr.Blocks(title="Chat with CSV using Llama2 🦙") as interface:
-        # Create unique session ID for each user
         session_id = gr.State(lambda: str(uuid.uuid4()))
-        # Create user-specific chatbot instance
         chatbot_state = gr.State(lambda: None)
         
         gr.HTML("<h1 style='text-align: center;'>Chat with CSV using Llama2 🦙</h1>")
@@ -157,11 +158,11 @@ def create_gradio_interface():
                 
                 with gr.Accordion("Informasi Model", open=False):
                     gr.Markdown("""
-                    **Model**: Llama-2-7b-chat
+                    **Model**: Llama-2-7b-chat-hf
                     
                     **Fitur**:
                     - Dioptimalkan untuk analisis data dan percakapan
-                    - Efisien dengan kuantisasi GGUF
+                    - Menggunakan API Hugging Face untuk efisiensi
                     - Manajemen sesi per pengguna
                     """)
             
@@ -178,9 +179,8 @@ def create_gradio_interface():
                 submit_button = gr.Button("Kirim")
                 clear_button = gr.Button("Bersihkan Chat")
         
-        # Process file handler
+        # Handler functions
         def handle_process_file(file, sess_id):
-            # Create chatbot if doesn't exist
             chatbot = ChatBot(sess_id)
             result = chatbot.process_file(file)
             return chatbot, [(None, result)]
@@ -191,14 +191,11 @@ def create_gradio_interface():
             outputs=[chatbot_state, chatbot_interface]
         )
         
-        # Chat handler - show user message immediately and then start thinking
         def user_message_submitted(message, history, chatbot, sess_id):
-            # Add user message to history immediately
             history = history + [(message, None)]
             return history, "", chatbot, sess_id
         
         def bot_response(history, chatbot, sess_id):
-            # Create chatbot if doesn't exist
             if chatbot is None:
                 chatbot = ChatBot(sess_id)
                 history[-1] = (history[-1][0], "Mohon upload file CSV terlebih dahulu.")
@@ -206,8 +203,6 @@ def create_gradio_interface():
             
             user_message = history[-1][0]
             response = chatbot.chat(user_message, history[:-1])
-            
-            # Update the last history item with the response
             history[-1] = (user_message, response)
             return chatbot, history
         
@@ -221,7 +216,6 @@ def create_gradio_interface():
             outputs=[chatbot_state, chatbot_interface]
         )
         
-        # Also hook up message input for pressing Enter
         message_input.submit(
             fn=user_message_submitted,
             inputs=[message_input, chatbot_interface, chatbot_state, session_id],
@@ -232,7 +226,6 @@ def create_gradio_interface():
             outputs=[chatbot_state, chatbot_interface]
         )
         
-        # Clear chat handler
         def handle_clear_chat(chatbot):
             if chatbot is not None:
                 chatbot.chat_history = []