Update app.py

app.py

@@ -17,6 +17,10 @@ dtype = None
 load_in_4bit = True
 peft_model_name = "limitedonly41/mistral7b_v3_4_categories"
 
+# Global variables to cache model and tokenizer
+cached_model = None
+cached_tokenizer = None
+
 def translate_text(text: str) -> str:
     """Translate text to English"""
     try:
@@ -28,22 +32,29 @@ def translate_text(text: str) -> str:
         return text[:4990]
 
 @spaces.GPU
-def predict_inference(translated_text: str) -> str:
+def predict_inference(translated_text: str) -> Tuple[str, str]:
     """GPU-accelerated inference function"""
+    global cached_model, cached_tokenizer
+    
     try:
         if len(translated_text) < 150:
-            return 'Short'
-
-        from unsloth import FastLanguageModel
+            return 'Short', 'Text too short for classification'
 
-        # Load model INSIDE the GPU function
-        model, tokenizer = FastLanguageModel.from_pretrained(
-            model_name=peft_model_name,
-            max_seq_length=max_seq_length,
-            dtype=dtype,
-            load_in_4bit=load_in_4bit,
-        )
-        FastLanguageModel.for_inference(model)
+        # Load model only once (cache it globally)
+        if cached_model is None or cached_tokenizer is None:
+            print("Loading model for the first time...")
+            from unsloth import FastLanguageModel
+            
+            cached_model, cached_tokenizer = FastLanguageModel.from_pretrained(
+                model_name=peft_model_name,
+                max_seq_length=max_seq_length,
+                dtype=dtype,
+                load_in_4bit=load_in_4bit,
+            )
+            FastLanguageModel.for_inference(cached_model)
+            print("Model loaded and cached successfully!")
+        else:
+            print("Using cached model...")
 
         prompt = f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
 
@@ -58,24 +69,30 @@ Categorize the website into one of the 4 categories:\n\n1) OTHER\n2) NEWS/BLOG\n
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
         with autocast(device_type='cuda'):
-            inputs = tokenizer(prompt, return_tensors="pt").to(device)
-            outputs = model.generate(**inputs, max_new_tokens=16, use_cache=True)
-            ans = tokenizer.batch_decode(outputs)[0]
+            inputs = cached_tokenizer(prompt, return_tensors="pt").to(device)
+            outputs = cached_model.generate(**inputs, max_new_tokens=16, use_cache=True)
+            ans = cached_tokenizer.batch_decode(outputs)[0]
 
+        # Extract the raw prediction
         ans_pred = ans.split('### Response:')[1].split('<')[0].strip()
 
+        # Process the prediction into final category
         if 'OTHER' in ans_pred:
-            return 'OTHER'
+            final_prediction = 'OTHER'
         elif 'NEWS/BLOG' in ans_pred:
-            return 'NEWS/BLOG'
+            final_prediction = 'NEWS/BLOG'
         elif 'E-commerce' in ans_pred:
-            return 'E-commerce'
+            final_prediction = 'E-commerce'
+        elif 'COMPANIES' in ans_pred:
+            final_prediction = 'COMPANIES'
         else:
-            return 'ERROR'
+            final_prediction = 'ERROR'
+
+        return final_prediction, ans_pred
 
     except Exception as e:
         print(f"Inference error: {e}")
-        return 'ERROR'
+        return 'ERROR', f'Error: {str(e)[:200]}'
 
 def scrape_url_with_curl_cffi(url: str) -> Tuple[str, str]:
     """Scrape URL using curl_cffi for better compatibility"""
@@ -125,10 +142,10 @@ def scrape_url_with_curl_cffi(url: str) -> Tuple[str, str]:
     except Exception as e:
         return f"Scraping error: {str(e)[:200]}", ""
 
-def process_single_url(url: str, progress=gr.Progress()) -> Tuple[str, str]:
-    """Process a single URL and return both scraped text and prediction"""
+def process_single_url(url: str, progress=gr.Progress()) -> Tuple[str, str, str]:
+    """Process a single URL and return prediction, raw prediction, and scraped text"""
     if not url.strip():
-        return "Please provide a URL to process.", ""
+        return "Please provide a URL to process.", "", ""
     
     # Clean the URL
     url = url.strip()
@@ -142,10 +159,10 @@ def process_single_url(url: str, progress=gr.Progress()) -> Tuple[str, str]:
         status, scraped_text = scrape_url_with_curl_cffi(url)
         
         if status != "success":
-            return status, ""
+            return status, "", ""
             
         if len(scraped_text) < 50:
-            return "Error: Could not extract meaningful content from the website", scraped_text[:2000]
+            return "Error: Could not extract meaningful content from the website", "", scraped_text[:2000]
         
         # Limit text length for display
         scraped_display = scraped_text[:2000] + "..." if len(scraped_text) > 2000 else scraped_text
@@ -154,7 +171,7 @@ def process_single_url(url: str, progress=gr.Progress()) -> Tuple[str, str]:
         
         # Check if text is too short for classification
         if len(scraped_text) < 150:
-            return "Short", scraped_display
+            return "Short", "Text too short", scraped_display
             
         # Translate text
         translated = translate_text(scraped_text[:4990])
@@ -162,15 +179,15 @@ def process_single_url(url: str, progress=gr.Progress()) -> Tuple[str, str]:
         progress(0.7, desc="Classifying website...")
         
         # Get prediction using GPU
-        prediction = predict_inference(translated)
+        prediction, raw_prediction = predict_inference(translated)
         
         progress(1.0, desc="Complete!")
         
-        return prediction, scraped_display
+        return prediction, raw_prediction, scraped_display
         
     except Exception as e:
         error_msg = f"Error processing URL: {str(e)[:200]}"
-        return error_msg, ""
+        return error_msg, "", ""
 
 def create_interface():
     with gr.Blocks(title="Website Category Classifier", theme=gr.themes.Soft()) as interface:
@@ -178,7 +195,7 @@ def create_interface():
         <div style="text-align: center; margin-bottom: 20px;">
             <h1>🔍 Website Category Classifier</h1>
             <p style="font-size: 18px; color: #666;">
-                Classify websites into categories: <strong>OTHER</strong>, <strong>NEWS/BLOG</strong>, or <strong>E-commerce</strong>
+                Classify websites into categories: <strong>OTHER</strong>, <strong>NEWS/BLOG</strong>, <strong>E-commerce</strong>, <strong>COMPANIES</strong>
             </p>
         </div>
         """)
@@ -207,7 +224,9 @@ def create_interface():
                         ["https://amazon.com"], 
                         ["https://github.com"],
                         ["https://cnn.com"],
-                        ["https://shopify.com"]
+                        ["https://shopify.com"],
+                        ["https://apple.com"],
+                        ["https://microsoft.com"]
                     ],
                     inputs=[url_input],
                     label="📋 Try these examples:"
@@ -215,16 +234,23 @@ def create_interface():
 
             with gr.Column(scale=2):
                 prediction_output = gr.Textbox(
-                    label="🎯 Classification Result",
-                    lines=3,
+                    label="🎯 Final Classification Result",
+                    lines=2,
                     interactive=False,
                     info="The predicted category for this website"
                 )
                 
+                raw_prediction_output = gr.Textbox(
+                    label="🤖 Raw Model Prediction",
+                    lines=2,
+                    interactive=False,
+                    info="The exact text output from the AI model"
+                )
+                
                 scraped_output = gr.Textbox(
                     label="📄 Scraped Content Preview (first 2000 characters)",
-                    lines=20,
-                    max_lines=25,
+                    lines=15,
+                    max_lines=20,
                     interactive=False,
                     info="Raw text content extracted from the website"
                 )
@@ -242,6 +268,7 @@ def create_interface():
             <ul>
                 <li><strong>NEWS/BLOG:</strong> News websites, blogs, articles, journalism sites</li>
                 <li><strong>E-commerce:</strong> Online stores, shopping sites, marketplaces</li>
+                <li><strong>COMPANIES:</strong> Corporate websites, business pages, company information</li>
                 <li><strong>OTHER:</strong> All other types of websites (documentation, portfolios, etc.)</li>
             </ul>
         </div>
@@ -250,7 +277,7 @@ def create_interface():
         process_btn.click(
             fn=process_single_url,
             inputs=[url_input],
-            outputs=[prediction_output, scraped_output],
+            outputs=[prediction_output, raw_prediction_output, scraped_output],
             show_progress=True
         )