Update app.py

app.py

@@ -1,7 +1,9 @@
+
 import streamlit as st
 import uuid
 import sys
 import requests
+from peft import *
 import bitsandbytes as bnb
 import pandas as pd
 import torch
@@ -9,6 +11,12 @@ import torch.nn as nn
 import transformers
 from datasets import load_dataset
 from huggingface_hub import notebook_login
+from peft import (
+    LoraConfig,
+    PeftConfig,
+    get_peft_model,
+    prepare_model_for_kbit_training,
+)
 from transformers import (
     AutoConfig,
     AutoModelForCausalLM,
@@ -66,13 +74,11 @@ st.markdown("""
         </style>
         """, unsafe_allow_html=True)
 
-# Load the model and tokenizer from Hugging Face Hub
-model_name = "tiiuae/falcon-7b-instruct" 
-model = AutoModelForCausalLM.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-# Load the dataset
-dataset = load_dataset("nisaar/Lawyer_GPT_India")  
+# Load the model outside the handle_input() function
+with open('model_saved.pkl', 'rb') as f:
+    model = pickle.load(f)
+    if not isinstance(model, str):
+        st.error("The loaded model is not valid.")
 
 def write_top_bar():
     col1, col2, col3 = st.columns([1,10,2])
@@ -105,13 +111,8 @@ def handle_input():
     if len(chat_history) == MAX_HISTORY_LENGTH:
         chat_history = chat_history[:-1]
 
-    # Find the most similar example in the dataset
-    closest_example = find_closest_example(input, dataset)  # Implement your own logic to find the closest example
-
-    # Generate response using the model
-    inputs = tokenizer.encode(closest_example, return_tensors="pt")
-    outputs = model.generate(inputs)
-    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    prompt = input
+    answer = model  # Replace the predict() method with the model itself
 
     chat_history.append((input, answer))
 
@@ -121,12 +122,6 @@ def handle_input():
     })
     st.session_state.input = ""
 
-def find_closest_example(input, dataset):
-    # Implement your own logic to find the closest example in the dataset based on the user input
-    # You can use techniques like cosine similarity, semantic similarity, or any other approach that fits your dataset and requirements
-    # Return the closest example as a string
-    pass
-
 def write_user_message(md):
     col1, col2 = st.columns([1,12])
 
@@ -153,4 +148,4 @@ with st.container():
         write_chat_message(a, q)
 
 st.markdown('---')
-input = st.text_input("You are talking to an AI, ask any question.", key="input", on_change=handle_input)
+input = st.text_input("You are talking to an AI, ask any question.", key="input", on_change=handle_input)