Upload 9 files

app.py

@@ -8,7 +8,7 @@ import sys
 app = Flask(__name__, static_folder='static')
 CORS(app)
 
-MODEL_NAME = "KASHH-4/Gemma-finetuned"
+MODEL_NAME = "KASHH-4/mistral_fine-tuned_new"
 
 print("\n" + "="*80)
 print("🚀 LEGALDOCS AI - MODEL INITIALIZATION")
@@ -31,13 +31,12 @@ print(f"   - Model max length: {tokenizer.model_max_length}")
 print(f"   - Pad token: {tokenizer.pad_token}")
 
 print("Loading YOUR model weights...")
-# Optimized for 18GB RAM (Hugging Face Spaces) - use 4-bit quantization for speed
-# This will use ~10-12GB RAM and be MUCH faster than 8-bit
+# Optimized for 18GB RAM with 4-bit quantization
 quantization_config = BitsAndBytesConfig(
-    load_in_4bit=True,  # 4-bit is faster and uses less RAM than 8-bit
-    bnb_4bit_compute_dtype=torch.float16,  # Use FP16 for computations (faster)
-    bnb_4bit_quant_type="nf4",  # NormalFloat4 - better quality than standard 4-bit
-    bnb_4bit_use_double_quant=True,  # Double quantization for better performance
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_use_double_quant=True,
 )
 
 model = AutoModelForCausalLM.from_pretrained(
@@ -46,7 +45,7 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="auto",
     low_cpu_mem_usage=True,
     trust_remote_code=True,
-    torch_dtype=torch.float16,  # Use FP16 for faster inference
+    torch_dtype=torch.float16,
 )
 
 print("✅ Model loaded successfully!")
@@ -90,7 +89,7 @@ def generate():
             return jsonify({'error': 'Missing prompt in request body'}), 400
         
         prompt = data['prompt']
-        max_new_tokens = data.get('max_new_tokens', 256)
+        max_new_tokens = data.get('max_new_tokens', 400)
         temperature = data.get('temperature', 0.7)
         top_p = data.get('top_p', 0.9)
         
@@ -119,16 +118,18 @@ def generate():
         sys.stdout.flush()
         generation_start = time.time()
         
-        # Use greedy decoding to avoid temperature/sampling issues
+        # Use controlled sampling for better JSON generation
         with torch.no_grad():
+            torch.set_num_threads(2)  # Use both CPU cores
             outputs = model.generate(
                 **inputs,
-                max_new_tokens=max_new_tokens,
-                do_sample=False,  # Use greedy decoding instead of sampling
+                max_new_tokens=400,
+                temperature=0.7,
+                top_p=0.9,
+                do_sample=True,
                 pad_token_id=tokenizer.eos_token_id,
                 eos_token_id=tokenizer.eos_token_id,
-                repetition_penalty=1.1,
-                no_repeat_ngram_size=3
+                repetition_penalty=1.1
             )
         
         generation_time = time.time() - generation_start