Update app.py

app.py

@@ -5,7 +5,8 @@ import PyPDF2
 from docx import Document
 import pandas as pd
 from dotenv import load_dotenv
-from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
+from unsloth import FastLanguageModel
+from transformers import AutoTokenizer
 
 # Load environment variables
 load_dotenv()
@@ -48,7 +49,7 @@ st.title("🚀 Ataliba o Agent Nerdx 🚀")
 # Sidebar
 with st.sidebar:
     st.header("⚡️ Hugging Face Model Loaded")
-    st.markdown("Model: amiguel/unsloth_finetune_test")
+    st.markdown("Model: `amiguel/unsloth_finetune_test` with LoRA")
     uploaded_file = st.file_uploader("Upload technical documents", type=["pdf", "docx", "xlsx", "xlsm"])
 
 # Session state
@@ -81,17 +82,23 @@ if uploaded_file and not st.session_state.file_context:
 
 # Load model
 @st.cache_resource
-def load_custom_model():
-    model_name = "amiguel/unsloth_finetune_test"
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForSequenceClassification.from_pretrained(model_name)
-    return pipeline("text-classification", model=model, tokenizer=tokenizer)
+def load_unsloth_model():
+    base_model = "unsloth/llama-3-8b-Instruct-bnb-4bit"
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name=base_model,
+        max_seq_length=2048,
+        dtype=None,
+        load_in_4bit=True,
+        adapter_path="amiguel/unsloth_finetune_test"
+    )
+    FastLanguageModel.for_inference(model)
+    return model, tokenizer
 
 # Generate response
 def generate_response(prompt):
     bio_triggers = ['who are you', 'ataliba', 'yourself', 'skilled at', 
                     'background', 'experience', 'valonylabs', 'totalenergies']
-    
+
     if any(trigger in prompt.lower() for trigger in bio_triggers):
         for line in ATALIBA_BIO.split('\n'):
             yield line + '\n'
@@ -99,16 +106,17 @@ def generate_response(prompt):
         return
 
     try:
-        classifier = load_custom_model()
-        result = classifier(prompt)[0]
-        label = result['label']
-        score = result['score']
-        context = st.session_state.file_context or "No document loaded."
-        
-        response_text = f"\n📘 **Prediction**: `{label}`\n📊 **Confidence**: `{score:.2%}`\n🗂️ **Context**: `{context[:300]}...`"
-        for line in response_text.split('\n'):
+        model, tokenizer = load_unsloth_model()
+        context = st.session_state.file_context or ""
+        full_prompt = f"You are an expert in Angolan labor law. Use the context to answer precisely.\nContext: {context}\n\nQuestion: {prompt}"
+
+        inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
+        outputs = model.generate(**inputs, max_new_tokens=256, do_sample=False)
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+        for line in response.split('\n'):
             yield line + '\n'
-            time.sleep(0.1)
+            time.sleep(0.05)
 
     except Exception as e:
         yield f"⚠️ Model Error: {str(e)}"