Update app.py

app.py

@@ -78,22 +78,46 @@ def predict(title, abstract):
     abstract = abstract.replace("\n", " ").strip().replace(''',"'")
     global model, tokenizer
     if model is None:
-        model = AutoModelForSequenceClassification.from_pretrained(
-            model_path,
-            num_labels=1,
-            torch_dtype=torch.float32 if device == 'cpu' else torch.float16
-        )
+        try:
+            # First try loading without quantization
+            model = AutoModelForSequenceClassification.from_pretrained(
+                model_path,
+                num_labels=1,
+                device_map='auto',
+                torch_dtype=torch.float32 if device == 'cpu' else torch.float16
+            )
+        except Exception as e:
+            print(f"Standard loading failed, trying without device mapping: {str(e)}")
+            # Fallback to basic loading
+            model = AutoModelForSequenceClassification.from_pretrained(
+                model_path,
+                num_labels=1,
+                torch_dtype=torch.float32
+            )
+            if torch.cuda.is_available():
+                model = model.cuda()
+                
         tokenizer = AutoTokenizer.from_pretrained(model_path)
-        model.to(device)
         model.eval()
+        
     text = f'''Given a certain paper, Title: {title}\n Abstract: {abstract}. \n Predict its normalized academic impact (between 0 and 1):'''
-    inputs = tokenizer(text, return_tensors="pt").to(device)
-    with torch.no_grad():
-        outputs = model(**inputs)
-    probability = torch.sigmoid(outputs.logits).item()
-    if probability + 0.05 >= 1.0:
-        return round(1, 4)
-    return round(probability + 0.05, 4)
+    
+    try:
+        inputs = tokenizer(text, return_tensors="pt")
+        if torch.cuda.is_available():
+            inputs = {k: v.cuda() for k, v in inputs.items()}
+            
+        with torch.no_grad():
+            outputs = model(**inputs)
+        probability = torch.sigmoid(outputs.logits).item()
+        
+        if probability + 0.05 >= 1.0:
+            return round(1, 4)
+        return round(probability + 0.05, 4)
+        
+    except Exception as e:
+        print(f"Prediction error: {str(e)}")
+        return 0.0  # Return default value in case of error
 
 def get_grade_and_emoji(score):
     if score >= 0.900: return "AAA 🌟"
@@ -224,8 +248,7 @@ css = """
 with gr.Blocks(theme=gr.themes.Default(), css=css) as iface:
     gr.Markdown(
         """
-        # PaperImpact: AI-Powered Research Impact Predictor {.main-title}
-        ### Estimate the future academic impact from the title and abstract with advanced AI analysis {.sub-title}
+        # PaperImpact: AI-Powered Research Impact Predictor 
         """
     )