app para avalição do modelo treinado

app.py

@@ -1,21 +1,35 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import torch
 
 # Load model and tokenizer
-model_name = "mlabonne/FineLlama-3.1-8B"  # Replace with your model's name
+model_name = "unsloth/Llama-3.2-1B-Instruct-bnb-4bit"  # Replace with your model's name
+
+# Configure quantization
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
+
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    quantization_config=bnb_config,
+    device_map="auto",
+)
 
-def generate_text(prompt, max_length, temperature):
+def generate_text(prompt, max_new_tokens, temperature):
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_length=max_length,
+            max_new_tokens=max_new_tokens,
             temperature=temperature,
-            num_return_sequences=1
+            num_return_sequences=1,
+            do_sample=True,
         )
     
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
@@ -25,7 +39,7 @@ iface = gr.Interface(
     fn=generate_text,
     inputs=[
         gr.Textbox(lines=5, label="Enter your ESG-related prompt"),
-        gr.Slider(50, 500, value=200, label="Maximum Length"),
+        gr.Slider(50, 500, value=200, label="Maximum New Tokens"),
         gr.Slider(0.1, 2.0, value=0.7, step=0.1, label="Temperature")
     ],
     outputs=gr.Textbox(label="Generated ESG Report Paragraph"),

requirements.txt

@@ -2,4 +2,6 @@ huggingface_hub==0.25.2
 
 gradio
 transformers
-torch
+torch
+accelerate>=0.26.0
+bitsandbytes