Update app.py

app.py

@@ -18,7 +18,7 @@ tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 try:
     model = AutoModel.from_pretrained(
         MODEL_ID,
-        torch_dtype=torch.bfloat16,
+        dtype=torch.bfloat16,  # 'torch_dtype' warning fixed here
         device_map="auto",
         trust_remote_code=True
     )
@@ -28,7 +28,7 @@ except KeyError:
     model = AutoModel.from_pretrained(
         MODEL_ID,
         config=config,
-        torch_dtype=torch.bfloat16,
+        dtype=torch.bfloat16,
         device_map="auto",
         trust_remote_code=True
     )
@@ -44,7 +44,7 @@ def custom_generate(user_input):
     lm_head = None
     with torch.no_grad():
         out = model(input_ids)
-        h = out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0]
+        h = out.last_hidden_state if hasattr(out, "last_hidden_state") else out
         dim = h.shape[-1]
         
         for name, module in model.named_modules():
@@ -60,7 +60,7 @@ def custom_generate(user_input):
     if not lm_head:
         return "Error: Language head not found."
 
-    generated_ids = input_ids[0].tolist()
+    generated_ids = input_ids[0].tolist()  # Fixed slicing error here
     start_len = len(generated_ids)
     
     # Generate tokens
@@ -68,7 +68,7 @@ def custom_generate(user_input):
         curr_tensor = torch.tensor([generated_ids]).to(model.device)
         with torch.no_grad():
             out = model(curr_tensor)
-            h = out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0]
+            h = out.last_hidden_state if hasattr(out, "last_hidden_state") else out
             logits = lm_head(h[:, -1, :])
             
             token = torch.argmax(logits, dim=-1).item()
@@ -78,19 +78,19 @@ def custom_generate(user_input):
                 break
                 
     response = tokenizer.decode(generated_ids[start_len:], skip_special_tokens=True)
-    return response.split("User:")[0].strip()
+    return response.split("User:")[0].strip()  # Fixed string parsing logic here
 
 # Gradio Interface Chat Setup
 def chat_interface(message, history):
     return custom_generate(message)
 
-# Launching Web UI
-demo = gr.ChatInterface(
-    fn=chat_interface,
-    title="BOL-AI v1.0 PRO",
-    description="Developer: Vivek Vijay Dalvi | Company: MAHAVEER AI",
-    theme="soft"
-)
+# Launching Web UI with explicit theme block to fix TypeError
+with gr.Blocks(theme="soft") as demo:
+    gr.ChatInterface(
+        fn=chat_interface,
+        title="BOL-AI v1.0 PRO",
+        description="Developer: Vivek Vijay Dalvi | Company: MAHAVEER AI"
+    )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()