Update app.py

app.py

@@ -22,7 +22,7 @@ from transformers import (
     AutoTokenizer,
     BitsAndBytesConfig,
 )
-
+import pickle
 
 USER_ICON = "images/user-icon.png"
 AI_ICON = "images/ai-icon.png"
@@ -104,123 +104,15 @@ def handle_input():
     if len(chat_history) == MAX_HISTORY_LENGTH:
         chat_history = chat_history[:-1]
 
-    # api_url = "https://9pl792yjf9.execute-api.us-east-1.amazonaws.com/beta/chatcogwise"
-    # api_request_data = {"question": input, "session": user_id}
-    # api_response = requests.post(api_url, json=api_request_data)
-    # result = api_response.json()
-
-    # answer = result['answer']
-    # !pip install -Uqqq pip --progress-bar off
-    # !pip install -qqq bitsandbytes == 0.39.0
-    # !pip install -qqqtorch --2.0.1 --progress-bar off
-    # !pip install -qqq -U git + https://github.com/huggingface/transformers.git@e03a9cc --progress-bar off
-    # !pip install -qqq -U git + https://github.com/huggingface/peft.git@42a184f --progress-bar off
-    # !pip install -qqq -U git + https://github.com/huggingface/accelerate.git@c9fbb71 --progress-bar off
-    # !pip install -qqq datasets == 2.12.0 --progress-bar off
-    # !pip install -qqq loralib == 0.1.1 --progress-bar off
-    # !pip install einops
-
-    import os
-    # from pprint import pprint
-    # import json
-
-
-
-    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
-
-    # notebook_login()
-    # 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
-
-    # """### Load dataset"""
-
-    from datasets import load_dataset
-
-    dataset_name = "nisaar/Lawyer_GPT_India"
-    # dataset_name = "patrick11434/TEST_LLM_DATASET"
-    dataset = load_dataset(dataset_name, split="train")
-
-    # """## Load adapters from the Hub
-
-    # You can also directly load adapters from the Hub using the commands below:
-    # """
-
-
-    # change peft_model_id
-    bnb_config = BitsAndBytesConfig(
-        load_in_4bit=True,
-        load_4bit_use_double_quant=True,
-        bnb_4bit_quant_type="nf4",
-        bnb_4bit_compute_dtype=torch.bfloat16,
-    )
-
-    peft_model_id = "nisaar/falcon7b-Indian_Law_150Prompts"
-    config = PeftConfig.from_pretrained(peft_model_id)
-    model = AutoModelForCausalLM.from_pretrained(
-        config.base_model_name_or_path,
-        return_dict=True,
-        quantization_config=bnb_config,
-        device_map="auto",
-        trust_remote_code=True,
-    )
-    tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
-    tokenizer.pad_token = tokenizer.eos_token
-
-    model = PeftModel.from_pretrained(model, peft_model_id)
-
-    """## Inference
-
-    You can then directly use the trained model or the model that you have loaded from the 🤗 Hub for inference as you would do it usually in `transformers`.
-    """
-
-    generation_config = model.generation_config
-    generation_config.max_new_tokens = 200
-    generation_config_temperature = 1
-    generation_config.top_p = 0.7
-    generation_config.num_return_sequences = 1
-    generation_config.pad_token_id = tokenizer.eos_token_id
-    generation_config_eod_token_id = tokenizer.eos_token_id
-
-    DEVICE = "cuda:0"
-
-    # Commented out IPython magic to ensure Python compatibility.
-    # %%time
-    # prompt = f"""
-    # <human>: Who appoints the Chief Justice of India?
-    # <assistant>:
-    # """.strip()
-    #
-    # encoding = tokenizer(prompt, return_tensors="pt").to(DEVICE)
-    # with torch.inference_mode():
-    #   outputs = model.generate(
-    #       input_ids=encoding.attention_mask,
-    #       generation_config=generation_config,
-    #   )
-    # print(tokenizer.decode(outputs[0],skip_special_tokens=True))
+    with open('model_saved.pkl', 'rb') as f:
+        model = pickle.load(f)
 
     def generate_response(question: str) -> str:
-        prompt = f"""
-        <human>: {question}
-        <assistant>:
-        """.strip()
-        encoding = tokenizer(prompt, return_tensors="pt").to(DEVICE)
-        with torch.inference_mode():
-            outputs = model.generate(
-                input_ids=encoding.input_ids,
-                attention_mask=encoding.attention_mask,
-                generation_config=generation_config,
-            )
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-        assistant_start = '<assistant>:'
-        response_start = response.find(assistant_start)
-        return response[response_start + len(assistant_start):].strip()
+        return model.predict([question])[0]
 
-    # prompt = "Debate the merits and demerits of introducing simultaneous elections in India?"
-    prompt=input
-    answer=generate_response(prompt)
-    print(answer)
+    prompt = input
+    answer = generate_response(prompt)
 
-    # answer='Yes'
     chat_history.append((input, answer))
 
     st.session_state.answers.append({