Update app.py

Fixed some changes

app.py

@@ -1,42 +1,48 @@
 import torch
-from trl import SFTTrainer
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    BitsAndBytesConfig,
+    pipeline
+)
 from peft import LoraConfig
-from datasets import load_dataset
-from transformers import (AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TrainingArguments, pipeline)
-import ipywidgets as widgets
-from IPython.display import display
 import gradio as gr
 
-llama_model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path = "aboonaji/llama2finetune-v2",
-                                                   quantization_config = BitsAndBytesConfig(load_in_4bit = True, bnb_4bit_compute_dtype = getattr(torch, "float16"), bnb_4bit_quant_type = "nf4"))
+# Load pre-trained model
+llama_model = AutoModelForCausalLM.from_pretrained(
+    pretrained_model_name_or_path="aboonaji/llama2finetune-v2",
+    quantization_config=BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_compute_dtype=getattr(torch, "float16"),
+        bnb_4bit_quant_type="nf4"
+    ),
+    device_map="auto"
+)
 llama_model.config.use_cache = False
 llama_model.config.pretraining_tp = 1
 
-llama_tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path = "aboonaji/llama2finetune-v2", trust_remote_code = True)
+llama_tokenizer = AutoTokenizer.from_pretrained(
+    pretrained_model_name_or_path="aboonaji/llama2finetune-v2",
+    trust_remote_code=True
+)
 llama_tokenizer.pad_token = llama_tokenizer.eos_token
 llama_tokenizer.padding_side = "right"
 
-training_arguments = TrainingArguments(output_dir = "./results", per_device_train_batch_size = 1, max_steps = 100)
-llama_sft_trainer = SFTTrainer(model = llama_model,
-                               args = training_arguments,
-                               train_dataset = load_dataset(path = "aboonaji/wiki_medical_terms_llam2_format", split = "train"),
-                               tokenizer = llama_tokenizer,
-                               peft_config = LoraConfig(task_type = "CAUSAL_LM", r = 16, lora_alpha = 16, lora_dropout = 0.1),
-                               dataset_text_field = "text")
-
-llama_sft_trainer.train()
-
-generator = pipeline("text-generation", model=llama_model, tokenizer=llama_tokenizer, max_length=500)
+# Pipeline for inference
+generator = pipeline(
+    "text-generation",
+    model=llama_model,
+    tokenizer=llama_tokenizer,
+    max_length=500
+)
 
 # In-memory user database
 user_db = {}
 
-# Response function
 def generate_response(prompt):
     response = generator(f"<s>[INST] {prompt} [/INST]")[0]["generated_text"]
     return response
 
-# Sign-up and login logic
 def signup_user(new_username, new_password):
     if not new_username or not new_password:
         return "❌ No input. Please provide both username and password."
@@ -48,10 +54,10 @@ def signup_user(new_username, new_password):
 def login_user(username, password):
     if username in user_db and user_db[username] == password:
         return (
-            gr.update(visible=True),   # Show chat UI
-            gr.update(visible=False),  # Hide login UI
-            "",                         # Clear login message
-            gr.update(selected=3)      # Switch to Chat tab
+            gr.update(visible=True),
+            gr.update(visible=False),
+            "",
+            gr.update(selected=3)
         )
     return (
         gr.update(visible=False),
@@ -62,9 +68,9 @@ def login_user(username, password):
 
 def logout_user():
     return (
-        gr.update(visible=False),  # Hide chat UI
-        gr.update(visible=True),   # Show login UI
-        gr.update(selected=0)      # Switch to Landing tab
+        gr.update(visible=False),
+        gr.update(visible=True),
+        gr.update(selected=0)
     )
 
 with gr.Blocks(theme="soft", css="""
@@ -130,23 +136,16 @@ with gr.Blocks(theme="soft", css="""
                     inputs=prompt
                 )
 
-
-
                 logout_btn.click(fn=logout_user, outputs=[chat_ui, login_ui, tabs])
 
-    # Tab navigation logic
+    # Navigation and interaction logic
     to_signup = gr.Button(visible=False)
     to_login = gr.Button(visible=False)
 
     to_signup.click(lambda: gr.update(selected=1), outputs=tabs)
     to_login.click(lambda: gr.update(selected=2), outputs=tabs)
     create_account_btn.click(fn=signup_user, inputs=[new_username, new_password], outputs=signup_msg)
-    login_btn.click(
-        fn=login_user,
-        inputs=[username, password],
-        outputs=[chat_ui, login_ui, login_msg, tabs]
-    )
+    login_btn.click(fn=login_user, inputs=[username, password], outputs=[chat_ui, login_ui, login_msg, tabs])
     submit_btn.click(fn=generate_response, inputs=prompt, outputs=response)
 
-# Launch the interface
-demo.launch(share=True)
+demo.launch()