Update app.py

app.py

@@ -1,20 +1,20 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import gradio as gr
+import torch
 
 model_id = "AmiyendraOP/llama3-legal-finetuned"
 
-# Load tokenizer and model
+# Load tokenizer and model without quantization or bitsandbytes
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id, low_cpu_mem_usage=True, torch_dtype=torch.float32)
 
-# Set device properly
+# Set device
 device = 0 if torch.cuda.is_available() else -1
 
-# Use pipeline for text generation
+# Use the text-generation pipeline
 pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)
 
-# Define chat function
+# Define a chat function
 def chat(prompt):
     response = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7)[0]["generated_text"]
     return response
@@ -24,5 +24,5 @@ gr.Interface(
     fn=chat,
     inputs=gr.Textbox(lines=4, placeholder="Enter legal question...", label="Your Question"),
     outputs=gr.Textbox(label="Response"),
-    title="LLaMA 3 Legal Chatbot (Fine-tuned)",
+    title="LLaMA 3 Legal Chatbot (Fine-tuned)"
 ).launch()