Update app.py

app.py

@@ -1,9 +1,7 @@
-
 import streamlit as st
 import uuid
 import sys
 import requests
-from peft import *
 import bitsandbytes as bnb
 import pandas as pd
 import torch
@@ -11,12 +9,6 @@ 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,
@@ -74,11 +66,13 @@ st.markdown("""
         </style>
         """, unsafe_allow_html=True)
 
-# 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.")
+# 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")  
 
 def write_top_bar():
     col1, col2, col3 = st.columns([1,10,2])
@@ -111,8 +105,13 @@ def handle_input():
     if len(chat_history) == MAX_HISTORY_LENGTH:
         chat_history = chat_history[:-1]
 
-    prompt = input
-    answer = model  # Replace the predict() method with the model itself
+    # 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)
 
     chat_history.append((input, answer))
 
@@ -122,6 +121,12 @@ 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])
 
@@ -148,4 +153,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)