Update app.py

app.py

@@ -1,527 +1,236 @@
-import gradio as gr                                                                                                         
-import torch                                                                                                                
-from setfit import SetFitModel                                                                                              
-from transformers import AutoTokenizer, T5ForConditionalGeneration                                                          
-import json                                                                                                                 
-import logging                                                                                                              
-import re                                                                                                                   
-from typing import List, Dict, Any                                                                                          
-import os                                                                                                                   
-                                                                                                                          
-logging.basicConfig(level=logging.INFO)                                                                                     
-logger = logging.getLogger(__name__)                                                                                        
-                                                                                                                          
-# Global model variables                                                                                                    
-classifier_model = None                                                                                                     
-extractor_model = None                                                                                                      
-extractor_tokenizer = None                                                                                                  
-device = None                                                                                                               
-                                                                                                                          
-def load_models():                                                                                                          
-  """Load both classification and extraction models"""                                                                    
-  global classifier_model, extractor_model, extractor_tokenizer, device                                                   
-                                                                                                                          
-  # Set device (GPU if available)                                                                                         
-  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")                                                   
-  logger.info(f"🖥️ Using device: {device}")                                                                               
-                                                                                                                          
-  try:                                                                                                                    
-      # Load classifier                                                                                                   
-      classifier_name = "Tomiwajin/testClasifier"                                                                         
-      token = os.getenv("HF_TOKEN")                                                                                       
-                                                                                                                          
-      classifier_model = SetFitModel.from_pretrained(                                                                     
-          classifier_name,                                                                                                
-          use_auth_token=token if token else False                                                                        
-      )                                                                                                                   
-      logger.info(f"✅ Classifier loaded: {classifier_name}")                                                             
-                                                                                                                          
-      # Load extractor                                                                                                    
-      extractor_name = "Tomiwajin/email-company-role-extractor"                                                           
-      extractor_tokenizer = AutoTokenizer.from_pretrained(extractor_name)                                                 
-      extractor_model = T5ForConditionalGeneration.from_pretrained(extractor_name)                                        
-      extractor_model.to(device)                                                                                          
-      extractor_model.eval()                                                                                              
-      logger.info(f"✅ Extractor loaded: {extractor_name}")                                                               
-                                                                                                                          
-      return True                                                                                                         
-  except Exception as e:                                                                                                  
-      logger.error(f"❌ Model loading failed: {e}")                                                                       
-      return False                                                                                                        
-                                                                                                                          
-def classify_single_email(email_text: str) -> Dict[str, Any]:                                                               
-  """Classify a single email"""                                                                                           
-  if not classifier_model:                                                                                                
-      return {"error": "Classifier not loaded", "success": False}                                                         
-                                                                                                                          
-  try:                                                                                                                    
-      email_text = email_text.strip()[:1000]                                                                              
-      predictions = classifier_model.predict([email_text])                                                                
-      probabilities = classifier_model.predict_proba([email_text])[0]                                                     
-                                                                                                                          
-      predicted_label = predictions[0]                                                                                    
-      confidence = max(probabilities)                                                                                     
-                                                                                                                          
-      return {                                                                                                            
-          "label": str(predicted_label),                                                                                  
-          "score": round(float(confidence), 4),                                                                           
-          "success": True                                                                                                 
-      }                                                                                                                   
-  except Exception as e:                                                                                                  
-      logger.error(f"Classification error: {e}")                                                                          
-      return {"error": str(e), "success": False}                                                                          
-                                                                                                                          
-def extract_job_info(email_text: str) -> Dict[str, Any]:                                                                    
-  """Extract company and role from email"""                                                                               
-  if not extractor_model or not extractor_tokenizer:                                                                      
-      return {"error": "Extractor not loaded", "success": False}                                                          
-                                                                                                                          
-  try:                                                                                                                    
-      email_text = email_text.strip()[:1000]                                                                              
-      input_text = f"extract company and role: {email_text}"                                                              
-                                                                                                                          
-      inputs = extractor_tokenizer(                                                                                       
-          input_text,                                                                                                     
-          return_tensors='pt',                                                                                            
-          max_length=512,                                                                                                 
-          truncation=True                                                                                                 
-      ).to(device)                                                                                                        
-                                                                                                                          
-      with torch.no_grad():                                                                                               
-          outputs = extractor_model.generate(                                                                             
-              inputs.input_ids,                                                                                           
-              attention_mask=inputs.attention_mask,                                                                       
-              max_length=128,                                                                                             
-              num_beams=2,                                                                                                
-              early_stopping=True,                                                                                        
-              pad_token_id=extractor_tokenizer.pad_token_id                                                               
-          )                                                                                                               
-                                                                                                                          
-      prediction = extractor_tokenizer.decode(outputs[0], skip_special_tokens=True)                                       
-      result = parse_extraction_result(prediction)                                                                        
-      return result                                                                                                       
-  except Exception as e:                                                                                                  
-      logger.error(f"Extraction error: {e}")                                                                              
-      return {                                                                                                            
-          "company": "unknown",                                                                                           
-          "role": "unknown",                                                                                              
-          "success": False,                                                                                               
-          "error": str(e)                                                                                                 
-      }                                                                                                                   
-                                                                                                                          
-def parse_extraction_result(prediction: str) -> Dict[str, Any]:                                                             
-  """Parse the model output into structured result"""                                                                     
-  try:                                                                                                                    
-      fixed = prediction.strip()                                                                                          
-      if fixed.startswith('"') and not fixed.startswith('{'):                                                             
-          fixed = '{' + fixed                                                                                             
-      if not fixed.endswith('}'):                                                                                         
-          fixed = fixed + '}'                                                                                             
-      fixed = re.sub(r'",(\s*)"', '", "', fixed)                                                                          
-                                                                                                                          
-      result = json.loads(fixed)                                                                                          
-      return {                                                                                                            
-          "company": result.get("company", "unknown"),                                                                    
-          "role": result.get("role", "unknown"),                                                                          
-          "success": True                                                                                                 
-      }                                                                                                                   
-  except:                                                                                                                 
-      return {                                                                                                            
-          "company": "unknown",                                                                                           
-          "role": "unknown",                                                                                              
-          "success": False                                                                                                
-      }                                                                                                                   
-                                                                                                                          
-def classify_batch_emails(emails: List[str]) -> List[Dict[str, Any]]:                                                       
-  """Classify multiple emails - already batched"""                                                                        
-  if not classifier_model:                                                                                                
-      return [{"error": "Model not loaded", "success": False}] * len(emails)                                              
-                                                                                                                          
-  try:                                                                                                                    
-      cleaned_emails = [email.strip()[:1000] for email in emails]                                                         
-      predictions = classifier_model.predict(cleaned_emails)                                                              
-      probabilities = classifier_model.predict_proba(cleaned_emails)                                                      
-                                                                                                                          
-      results = []                                                                                                        
-      for pred, probs in zip(predictions, probabilities):                                                                 
-          results.append({                                                                                                
-              "label": str(pred),                                                                                         
-              "score": round(float(max(probs)), 4),                                                                       
-              "success": True                                                                                             
-          })                                                                                                              
-                                                                                                                          
-      return results                                                                                                      
-  except Exception as e:                                                                                                  
-      logger.error(f"Batch classification error: {e}")                                                                    
-      return [{"error": str(e), "success": False}] * len(emails)                                                          
-                                                                                                                          
-def extract_batch(emails: List[str]) -> List[Dict[str, Any]]:                                                               
-  """Extract company and role from multiple emails - BATCHED for performance"""                                           
-  if not extractor_model or not extractor_tokenizer:                                                                      
-      return [{"error": "Extractor not loaded", "success": False}] * len(emails)                                          
-                                                                                                                          
-  if len(emails) == 0:                                                                                                    
-      return []                                                                                                           
-                                                                                                                          
-  try:                                                                                                                    
-      # Prepare all inputs at once                                                                                        
-      cleaned_emails = [email.strip()[:1000] for email in emails]                                                         
-      input_texts = [f"extract company and role: {email}" for email in cleaned_emails]                                    
-                                                                                                                          
-      # Batch tokenize                                                                                                    
-      inputs = extractor_tokenizer(                                                                                       
-          input_texts,                                                                                                    
-          return_tensors='pt',                                                                                            
-          max_length=512,                                                                                                 
-          truncation=True,                                                                                                
-          padding=True                                                                                                    
-      ).to(device)                                                                                                        
-                                                                                                                          
-      # Batch generate - process all at once                                                                              
-      with torch.no_grad():                                                                                               
-          outputs = extractor_model.generate(                                                                             
-              inputs.input_ids,                                                                                           
-              attention_mask=inputs.attention_mask,                                                                       
-              max_length=128,                                                                                             
-              num_beams=2,  # Reduced from 4 for speed                                                                    
-              early_stopping=True,                                                                                        
-              pad_token_id=extractor_tokenizer.pad_token_id                                                               
-          )                                                                                                               
-                                                                                                                          
-      # Decode all at once                                                                                                
-      predictions = extractor_tokenizer.batch_decode(outputs, skip_special_tokens=True)                                   
-                                                                                                                          
-      # Parse results                                                                                                     
-      results = [parse_extraction_result(pred) for pred in predictions]                                                   
-      return results                                                                                                      
-                                                                                                                          
-  except Exception as e:                                                                                                  
-      logger.error(f"Batch extraction error: {e}")                                                                        
-      return [{"company": "unknown", "role": "unknown", "success": False, "error": str(e)}] * len(emails)                 
-                                                                                                                          
-def process_batch(emails: List[str], job_labels: List[str] = None, threshold: float = 0.5) -> Dict[str, Any]:               
-  """                                                                                                                     
-  Combined endpoint: classify emails and extract job info in ONE call.                                                    
-  Only extracts from emails classified as job-related.                                                                    
-  """                                                                                                                     
-  if job_labels is None:                                                                                                  
-      job_labels = ["applied", "rejected", "interview", "next-phase", "offer"]                                            
-                                                                                                                          
-  # Step 1: Classify all emails                                                                                           
-  classifications = classify_batch_emails(emails)                                                                         
-                                                                                                                          
-  # Step 2: Find job-related emails                                                                                       
-  job_indices = []                                                                                                        
-  job_emails = []                                                                                                         
-  for i, (email, cls) in enumerate(zip(emails, classifications)):                                                         
-      if cls.get("success") and cls.get("label", "").lower() in job_labels and cls.get("score", 0) >= threshold:          
-          job_indices.append(i)                                                                                           
-          job_emails.append(email)                                                                                        
-                                                                                                                          
-  # Step 3: Extract only from job-related emails (batched)                                                                
-  extractions = extract_batch(job_emails) if job_emails else []                                                           
-                                                                                                                          
-  # Step 4: Combine results                                                                                               
-  results = []                                                                                                            
-  extraction_idx = 0                                                                                                      
-  for i, cls in enumerate(classifications):                                                                               
-      result = {                                                                                                          
-          "classification": cls,                                                                                          
-          "extraction": None                                                                                              
-      }                                                                                                                   
-      if i in job_indices:                                                                                                
-          result["extraction"] = extractions[extraction_idx]                                                              
-          extraction_idx += 1                                                                                             
-      results.append(result)                                                                                              
-                                                                                                                          
-  return {                                                                                                                
-      "results": results,                                                                                                 
-      "total": len(emails),                                                                                               
-      "job_related": len(job_emails)                                                                                      
-  }                                                                                                                       
-                                                                                                                          
-def gradio_classify(email_text: str) -> str:                                                                                
-  """Gradio interface for classification"""                                                                               
-  if not email_text.strip():                                                                                              
-      return "Please enter some email text to classify."                                                                  
-                                                                                                                          
-  result = classify_single_email(email_text)                                                                              
-                                                                                                                          
-  if result.get("success"):                                                                                               
-      return f"""                                                                                                         
-**Classification Result:**                                                                                                  
-- **Label:** {result['label']}                                                                                              
-- **Confidence:** {result['score']:.2%}                                                                                     
-"""                                                                                                                         
-  else:                                                                                                                   
-      return f"**Error:** {result.get('error', 'Unknown error')}"                                                         
-                                                                                                                          
-def gradio_extract(email_text: str) -> str:                                                                                 
-  """Gradio interface for extraction"""                                                                                   
-  if not email_text.strip():                                                                                              
-      return "Please enter some email text to extract from."                                                              
-                                                                                                                          
-  result = extract_job_info(email_text)                                                                                   
-                                                                                                                          
-  if result.get("success"):                                                                                               
-      return f"""                                                                                                         
-**Extraction Result:**                                                                                                      
-- **Company:** {result['company']}                                                                                          
-- **Role:** {result['role']}                                                                                                
-"""                                                                                                                         
-  else:                                                                                                                   
-      return f"**Error:** {result.get('error', 'Unknown error')}"                                                         
-                                                                                                                          
-def api_classify(email_text: str) -> Dict[str, Any]:                                                                        
-  """API endpoint for single classification"""                                                                            
-  return classify_single_email(email_text)                                                                                
-                                                                                                                          
-def api_extract(email_text: str) -> Dict[str, Any]:                                                                         
-  """API endpoint for single extraction"""                                                                                
-  return extract_job_info(email_text)                                                                                     
-                                                                                                                          
-def api_classify_batch(emails_json: str) -> str:                                                                            
-  """API endpoint for batch classification"""                                                                             
-  try:                                                                                                                    
-      emails = json.loads(emails_json)                                                                                    
-      if not isinstance(emails, list):                                                                                    
-          return json.dumps({"error": "Input must be a JSON array of strings"})                                           
-                                                                                                                          
-      if len(emails) > 400:                                                                                               
-          return json.dumps({"error": "Maximum 400 emails per batch"})                                                    
-                                                                                                                          
-      results = classify_batch_emails(emails)                                                                             
-      return json.dumps({"results": results})                                                                             
-  except json.JSONDecodeError:                                                                                            
-      return json.dumps({"error": "Invalid JSON format"})                                                                 
-  except Exception as e:                                                                                                  
-      return json.dumps({"error": str(e)})                                                                                
-                                                                                                                          
-def api_extract_batch(emails_json: str) -> str:                                                                             
-  """API endpoint for batch extraction - NOW BATCHED"""                                                                   
-  try:                                                                                                                    
-      emails = json.loads(emails_json)                                                                                    
-      if not isinstance(emails, list):                                                                                    
-          return json.dumps({"error": "Input must be a JSON array of strings"})                                           
-                                                                                                                          
-      if len(emails) > 400:                                                                                               
-          return json.dumps({"error": "Maximum 400 emails per batch"})                                                    
-                                                                                                                          
-      results = extract_batch(emails)                                                                                     
-      return json.dumps({"results": results})                                                                             
-  except json.JSONDecodeError:                                                                                            
-      return json.dumps({"error": "Invalid JSON format"})                                                                 
-  except Exception as e:                                                                                                  
-      return json.dumps({"error": str(e)})                                                                                
-                                                                                                                          
-def api_process_batch(emails_json: str, threshold: float = 0.5) -> str:                                                     
-  """API endpoint for combined classify + extract in ONE call"""                                                          
-  try:                                                                                                                    
-      emails = json.loads(emails_json)                                                                                    
-      if not isinstance(emails, list):                                                                                    
-          return json.dumps({"error": "Input must be a JSON array of strings"})                                           
-                                                                                                                          
-      if len(emails) > 400:                                                                                               
-          return json.dumps({"error": "Maximum 400 emails per batch"})                                                    
-                                                                                                                          
-      results = process_batch(emails, threshold=threshold)                                                                
-      return json.dumps(results)                                                                                          
-  except json.JSONDecodeError:                                                                                            
-      return json.dumps({"error": "Invalid JSON format"})                                                                 
-  except Exception as e:                                                                                                  
-      return json.dumps({"error": str(e)})                                                                                
-                                                                                                                          
-# Load models on startup                                                                                                    
-logger.info("Loading models...")                                                                                            
-models_loaded = load_models()                                                                                               
-                                                                                                                          
-if not models_loaded:                                                                                                       
-  logger.warning("Models failed to load - using dummy responses")                                                         
-                                                                                                                          
-# Create Gradio interface                                                                                                   
-with gr.Blocks(title="Email Classifier & Extractor", theme=gr.themes.Soft()) as demo:                                       
-  gr.Markdown("# 📧 Email Classification & Extraction API")                                                               
-  gr.Markdown("Classify job-related emails and extract company/role information.")                                        
-                                                                                                                          
-  with gr.Tab("Classification"):                                                                                          
-      with gr.Row():                                                                                                      
-          with gr.Column():                                                                                               
-              classify_input = gr.Textbox(                                                                                
-                  label="Email Content",                                                                                  
-                  placeholder="Paste your email content here...",                                                         
-                  lines=8,                                                                                                
-                  max_lines=20                                                                                            
-              )                                                                                                           
-              classify_btn = gr.Button("Classify Email", variant="primary")                                               
-                                                                                                                          
-          with gr.Column():                                                                                               
-              classify_output = gr.Markdown(label="Classification Result")                                                
-                                                                                                                          
-      classify_btn.click(                                                                                                 
-          fn=gradio_classify,                                                                                             
-          inputs=classify_input,                                                                                          
-          outputs=classify_output                                                                                         
-      )                                                                                                                   
-                                                                                                                          
-  with gr.Tab("Extraction"):                                                                                              
-      with gr.Row():                                                                                                      
-          with gr.Column():                                                                                               
-              extract_input = gr.Textbox(                                                                                 
-                  label="Email Content",                                                                                  
-                  placeholder="Paste job application email here...",                                                      
-                  lines=8,                                                                                                
-                  max_lines=20                                                                                            
-              )                                                                                                           
-              extract_btn = gr.Button("Extract Info", variant="primary")                                                  
-                                                                                                                          
-          with gr.Column():                                                                                               
-              extract_output = gr.Markdown(label="Extraction Result")                                                     
-                                                                                                                          
-      extract_btn.click(                                                                                                  
-          fn=gradio_extract,                                                                                              
-          inputs=extract_input,                                                                                           
-          outputs=extract_output                                                                                          
-      )                                                                                                                   
-                                                                                                                          
-  with gr.Tab("API Testing"):                                                                                             
-      gr.Markdown("### Test API Endpoints")                                                                               
-                                                                                                                          
-      with gr.Row():                                                                                                      
-          with gr.Column():                                                                                               
-              gr.Markdown("**Single Classification**")                                                                    
-              api_classify_input = gr.Textbox(label="Email Text", lines=4)                                                
-              api_classify_btn = gr.Button("Test Classify API")                                                           
-              api_classify_output = gr.JSON(label="Response")                                                             
-                                                                                                                          
-              api_classify_btn.click(                                                                                     
-                  fn=api_classify,                                                                                        
-                  inputs=api_classify_input,                                                                              
-                  outputs=api_classify_output,                                                                            
-                  api_name="classify"                                                                                     
-              )                                                                                                           
-                                                                                                                          
-          with gr.Column():                                                                                               
-              gr.Markdown("**Single Extraction**")                                                                        
-              api_extract_input = gr.Textbox(label="Email Text", lines=4)                                                 
-              api_extract_btn = gr.Button("Test Extract API")                                                             
-              api_extract_output = gr.JSON(label="Response")                                                              
-                                                                                                                          
-              api_extract_btn.click(                                                                                      
-                  fn=api_extract,                                                                                         
-                  inputs=api_extract_input,                                                                               
-                  outputs=api_extract_output,                                                                             
-                  api_name="extract"                                                                                      
-              )                                                                                                           
-                                                                                                                          
-      with gr.Row():                                                                                                      
-          with gr.Column():                                                                                               
-              gr.Markdown("**Batch Classification**")                                                                     
-              batch_classify_input = gr.Textbox(                                                                          
-                  label="JSON Array of Emails",                                                                           
-                  lines=6,                                                                                                
-                  placeholder='["Email 1", "Email 2"]'                                                                    
-              )                                                                                                           
-              batch_classify_btn = gr.Button("Test Batch Classify")                                                       
-              batch_classify_output = gr.Code(label="Response", language="json")                                          
-                                                                                                                          
-              batch_classify_btn.click(                                                                                   
-                  fn=api_classify_batch,                                                                                  
-                  inputs=batch_classify_input,                                                                            
-                  outputs=batch_classify_output,                                                                          
-                  api_name="classify_batch"                                                                               
-              )                                                                                                           
-                                                                                                                          
-          with gr.Column():                                                                                               
-              gr.Markdown("**Batch Extraction**")                                                                         
-              batch_extract_input = gr.Textbox(                                                                           
-                  label="JSON Array of Emails",                                                                           
-                  lines=6,                                                                                                
-                  placeholder='["Email 1", "Email 2"]'                                                                    
-              )                                                                                                           
-              batch_extract_btn = gr.Button("Test Batch Extract")                                                         
-              batch_extract_output = gr.Code(label="Response", language="json")                                           
-                                                                                                                          
-              batch_extract_btn.click(                                                                                    
-                  fn=api_extract_batch,                                                                                   
-                  inputs=batch_extract_input,                                                                             
-                  outputs=batch_extract_output,                                                                           
-                  api_name="extract_batch"                                                                                
-              )                                                                                                           
-                                                                                                                          
-      with gr.Row():                                                                                                      
-          with gr.Column(scale=2):                                                                                        
-              gr.Markdown("**🚀 Combined Process (Recommended)**")                                                        
-              gr.Markdown("*Classify + Extract in ONE call - fastest option*")                                            
-              process_input = gr.Textbox(                                                                                 
-                  label="JSON Array of Emails",                                                                           
-                  lines=6,                                                                                                
-                  placeholder='["Email 1", "Email 2"]'                                                                    
-              )                                                                                                           
-              process_threshold = gr.Slider(                                                                              
-                  minimum=0.1, maximum=0.9, value=0.5, step=0.1,                                                          
-                  label="Classification Threshold"                                                                        
-              )                                                                                                           
-              process_btn = gr.Button("Test Process API", variant="primary")                                              
-              process_output = gr.Code(label="Response", language="json")                                                 
-                                                                                                                          
-              process_btn.click(                                                                                          
-                  fn=api_process_batch,                                                                                   
-                  inputs=[process_input, process_threshold],                                                              
-                  outputs=process_output,                                                                                 
-                  api_name="process_batch"                                                                                
-              )                                                                                                           
-                                                                                                                          
-  with gr.Tab("Documentation"):                                                                                           
-      gr.Markdown(f"""                                                                                                    
-### Model Status                                                                                                            
-- **Status:** {'✅ Loaded' if models_loaded else '❌ Failed to load'}                                                       
-- **Device:** {device if device else 'Not initialized'}                                                                     
-- **Classifier:** SetFit (job categories)                                                                                   
-- **Extractor:** T5-small (company/role)                                                                                    
-                                                                                                                          
-### API Endpoints                                                                                                           
-                                                                                                                          
-#### 1. Single Classification                                                                                               
-`/api/classify` - Returns job category label                                                                                
-                                                                                                                          
-#### 2. Single Extraction                                                                                                   
-`/api/extract` - Returns company and role                                                                                   
-                                                                                                                          
-#### 3. Batch Classification                                                                                                
-`/api/classify_batch` - Classify multiple emails (max 400)                                                                  
-                                                                                                                          
-#### 4. Batch Extraction                                                                                                    
-`/api/extract_batch` - Extract from multiple emails (max 400) - **NOW BATCHED!**                                            
-                                                                                                                          
-#### 5. 🚀 Combined Process (NEW - FASTEST)                                                                                 
-`/api/process_batch` - Classify AND extract in ONE call                                                                     
-- Only extracts from job-related emails                                                                                     
-- Reduces API calls from 2 to 1                                                                                             
-- Recommended for best performance                                                                                          
-                                                                                                                          
-### Categories                                                                                                              
-- `applied` - Application submitted                                                                                         
-- `rejected` - Application rejected                                                                                         
-- `interview` - Interview invitation                                                                                        
-- `next-phase` - Next round invitation                                                                                      
-- `offer` - Job offer received                                                                                              
-- `other` - Not job-related                                                                                                 
-                                                                                                                          
-### Usage from Next.js                                                                                                      
-                                                                                                                          
-**Option 1: Separate calls (existing)**                                                                                     
-```javascript                                                                                                               
-const classifications = await client.predict("/classify_batch", {{                                                          
-emails_json: JSON.stringify(emails)                                                                                       
-}});                                                                                                                        
-                                                                                                                          
-const extractions = await client.predict("/extract_batch", {{                                                               
-emails_json: JSON.stringify(jobEmails)                                                                                    
-}});                                                                                                                        
-                                                                                                                          
-Option 2: Combined call (recommended - 2x faster)                                                                           
-const results = await client.predict("/process_batch", {{                                                                   
-emails_json: JSON.stringify(emails),                                                                                      
-threshold: 0.5                                                                                                            
-}});                                                                                                                        
-// Returns both classification AND extraction in one call     
+import gradio as gr
+import torch
+from setfit import SetFitModel
+from transformers import AutoTokenizer, T5ForConditionalGeneration
+import json
+import logging
+import re
+from typing import List, Dict, Any
+import os
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+classifier_model = None
+extractor_model = None
+extractor_tokenizer = None
+device = None
+
+def load_models():
+    global classifier_model, extractor_model, extractor_tokenizer, device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    logger.info(f"Using device: {device}")
+
+    try:
+        classifier_name = "Tomiwajin/testClasifier"
+        token = os.getenv("HF_TOKEN")
+        classifier_model = SetFitModel.from_pretrained(
+            classifier_name,
+            use_auth_token=token if token else False
+        )
+        logger.info(f"Classifier loaded: {classifier_name}")
+
+        extractor_name = "Tomiwajin/email-company-role-extractor"
+        extractor_tokenizer = AutoTokenizer.from_pretrained(extractor_name)
+        extractor_model = T5ForConditionalGeneration.from_pretrained(extractor_name)
+        extractor_model.to(device)
+        extractor_model.eval()
+        logger.info(f"Extractor loaded: {extractor_name}")
+        return True
+    except Exception as e:
+        logger.error(f"Model loading failed: {e}")
+        return False
+
+def parse_extraction_result(prediction):
+    try:
+        fixed = prediction.strip()
+        if fixed.startswith('"') and not fixed.startswith('{'):
+            fixed = '{' + fixed
+        if not fixed.endswith('}'):
+            fixed = fixed + '}'
+        fixed = re.sub(r'",(\s*)"', '", "', fixed)
+        result = json.loads(fixed)
+        return {
+            "company": result.get("company", "unknown"),
+            "role": result.get("role", "unknown"),
+            "success": True
+        }
+    except:
+        return {"company": "unknown", "role": "unknown", "success": False}
+
+def classify_single_email(email_text):
+    if not classifier_model:
+        return {"error": "Classifier not loaded", "success": False}
+    try:
+        email_text = email_text.strip()[:1000]
+        predictions = classifier_model.predict([email_text])
+        probabilities = classifier_model.predict_proba([email_text])[0]
+        return {
+            "label": str(predictions[0]),
+            "score": round(float(max(probabilities)), 4),
+            "success": True
+        }
+    except Exception as e:
+        logger.error(f"Classification error: {e}")
+        return {"error": str(e), "success": False}
+
+def extract_job_info(email_text):
+    if not extractor_model or not extractor_tokenizer:
+        return {"error": "Extractor not loaded", "success": False}
+    try:
+        email_text = email_text.strip()[:1000]
+        input_text = f"extract company and role: {email_text}"
+        inputs = extractor_tokenizer(
+            input_text, return_tensors='pt', max_length=512, truncation=True
+        ).to(device)
+        with torch.no_grad():
+            outputs = extractor_model.generate(
+                inputs.input_ids,
+                attention_mask=inputs.attention_mask,
+                max_length=128,
+                num_beams=2,
+                early_stopping=True,
+                pad_token_id=extractor_tokenizer.pad_token_id
+            )
+        prediction = extractor_tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return parse_extraction_result(prediction)
+    except Exception as e:
+        logger.error(f"Extraction error: {e}")
+        return {"company": "unknown", "role": "unknown", "success": False}
+
+def classify_batch_emails(emails):
+    if not classifier_model:
+        return [{"error": "Model not loaded", "success": False}] * len(emails)
+    try:
+        cleaned = [e.strip()[:1000] for e in emails]
+        predictions = classifier_model.predict(cleaned)
+        probabilities = classifier_model.predict_proba(cleaned)
+        return [
+            {"label": str(p), "score": round(float(max(pr)), 4), "success": True}
+            for p, pr in zip(predictions, probabilities)
+        ]
+    except Exception as e:
+        logger.error(f"Batch classification error: {e}")
+        return [{"error": str(e), "success": False}] * len(emails)
+
+def extract_batch(emails):
+    if not extractor_model or not extractor_tokenizer:
+        return [{"error": "Extractor not loaded", "success": False}] * len(emails)
+    if len(emails) == 0:
+        return []
+    try:
+        cleaned = [e.strip()[:1000] for e in emails]
+        input_texts = [f"extract company and role: {e}" for e in cleaned]
+        inputs = extractor_tokenizer(
+            input_texts, return_tensors='pt', max_length=512,
+            truncation=True, padding=True
+        ).to(device)
+        with torch.no_grad():
+            outputs = extractor_model.generate(
+                inputs.input_ids,
+                attention_mask=inputs.attention_mask,
+                max_length=128,
+                num_beams=2,
+                early_stopping=True,
+                pad_token_id=extractor_tokenizer.pad_token_id
+            )
+        predictions = extractor_tokenizer.batch_decode(outputs, skip_special_tokens=True)
+        return [parse_extraction_result(p) for p in predictions]
+    except Exception as e:
+        logger.error(f"Batch extraction error: {e}")
+        return [{"company": "unknown", "role": "unknown", "success": False}] * len(emails)
+
+def process_batch(emails, job_labels=None, threshold=0.5):
+    if job_labels is None:
+        job_labels = ["applied", "rejected", "interview", "next-phase", "offer"]
+    classifications = classify_batch_emails(emails)
+    job_indices = []
+    job_emails = []
+    for i, (email, cls) in enumerate(zip(emails, classifications)):
+        if cls.get("success") and cls.get("label", "").lower() in job_labels and cls.get("score", 0) >= threshold:
+            job_indices.append(i)
+            job_emails.append(email)
+    extractions = extract_batch(job_emails) if job_emails else []
+    results = []
+    ext_idx = 0
+    for i, cls in enumerate(classifications):
+        result = {"classification": cls, "extraction": None}
+        if i in job_indices:
+            result["extraction"] = extractions[ext_idx]
+            ext_idx += 1
+        results.append(result)
+    return {"results": results, "total": len(emails), "job_related": len(job_emails)}
+
+def api_classify_batch(emails_json):
+    try:
+        emails = json.loads(emails_json)
+        if not isinstance(emails, list):
+            return json.dumps({"error": "Input must be a JSON array"})
+        if len(emails) > 400:
+            return json.dumps({"error": "Maximum 400 emails per batch"})
+        results = classify_batch_emails(emails)
+        return json.dumps({"results": results})
+    except json.JSONDecodeError:
+        return json.dumps({"error": "Invalid JSON format"})
+    except Exception as e:
+        return json.dumps({"error": str(e)})
+
+def api_extract_batch(emails_json):
+    try:
+        emails = json.loads(emails_json)
+        if not isinstance(emails, list):
+            return json.dumps({"error": "Input must be a JSON array"})
+        if len(emails) > 400:
+            return json.dumps({"error": "Maximum 400 emails per batch"})
+        results = extract_batch(emails)
+        return json.dumps({"results": results})
+    except json.JSONDecodeError:
+        return json.dumps({"error": "Invalid JSON format"})
+    except Exception as e:
+        return json.dumps({"error": str(e)})
+
+def api_process_batch(emails_json, threshold=0.5):
+    try:
+        emails = json.loads(emails_json)
+        if not isinstance(emails, list):
+            return json.dumps({"error": "Input must be a JSON array"})
+        if len(emails) > 400:
+            return json.dumps({"error": "Maximum 400 emails per batch"})
+        results = process_batch(emails, threshold=threshold)
+        return json.dumps(results)
+    except json.JSONDecodeError:
+        return json.dumps({"error": "Invalid JSON format"})
+    except Exception as e:
+        return json.dumps({"error": str(e)})
+
+logger.info("Loading models...")
+models_loaded = load_models()
+
+with gr.Blocks(title="Email Classifier & Extractor", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# Email Classification & Extraction API")
+
+    with gr.Tab("Batch Classification"):
+        batch_input = gr.Textbox(label="JSON Array of Emails", lines=6, placeholder='["email1", "email2"]')
+        batch_btn = gr.Button("Classify Batch")
+        batch_output = gr.Code(label="Response", language="json")
+        batch_btn.click(fn=api_classify_batch, inputs=batch_input, outputs=batch_output, api_name="classify_batch")
+
+    with gr.Tab("Batch Extraction"):
+        extract_input = gr.Textbox(label="JSON Array of Emails", lines=6, placeholder='["email1", "email2"]')
+        extract_btn = gr.Button("Extract Batch")
+        extract_output = gr.Code(label="Response", language="json")
+        extract_btn.click(fn=api_extract_batch, inputs=extract_input, outputs=extract_output, api_name="extract_batch")
+
+    with gr.Tab("Combined Process"):
+        process_input = gr.Textbox(label="JSON Array of Emails", lines=6, placeholder='["email1", "email2"]')
+        process_threshold = gr.Slider(minimum=0.1, maximum=0.9, value=0.5, step=0.1, label="Threshold")
+        process_btn = gr.Button("Process Batch", variant="primary")
+        process_output = gr.Code(label="Response", language="json")
+        process_btn.click(fn=api_process_batch, inputs=[process_input, process_threshold], outputs=process_output, api_name="process_batch")
+
+    with gr.Tab("Status"):
+        status_text = "Loaded" if models_loaded else "Failed"
+        gr.Markdown(f"**Model Status:** {status_text}")
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860, show_api=True)