Update app.py

app.py

@@ -1,56 +1,59 @@
 import streamlit as st
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+import torch
 import json
-import requests
 import os
 
 # Streamlit app configuration
 st.set_page_config(page_title="AI Chatbot", layout="centered")
 
-# Fix the model's configuration before loading
+# Fix the model's configuration dynamically
 def fix_model_config(model_name):
-    # Download the configuration file from the model repository
+    # Load the configuration file directly from the Hugging Face hub
     config_url = f"https://huggingface.co/{model_name}/resolve/main/config.json"
     config_path = "config.json"
-    
+
     if not os.path.exists(config_path):
+        import requests
         response = requests.get(config_url)
-        response.raise_for_status()  # Raise an error if the request fails
+        response.raise_for_status()  # Ensure the request is successful
         with open(config_path, "w") as f:
             f.write(response.text)
-    
-    # Load the configuration and modify rope_scaling if necessary
+
+    # Load the configuration JSON
     with open(config_path, "r") as f:
         config = json.load(f)
-    
+
+    # Fix the `rope_scaling` field
     if "rope_scaling" in config:
         config["rope_scaling"] = {
-            "type": "linear",  # Replace the problematic rope_scaling type
+            "type": "linear",  # Replace the problematic structure with supported format
             "factor": config["rope_scaling"].get("factor", 1.0)
         }
-    
-    # Save the modified configuration
+
+    # Save the fixed configuration locally
     with open(config_path, "w") as f:
         json.dump(config, f)
-    
+
     return config_path
 
-# Load the model pipeline
+# Load the pipeline
 @st.cache_resource
 def load_pipeline():
     model_name = "Orenguteng/Llama-3.1-8B-Lexi-Uncensored-V2"
-    
+
     # Fix the model configuration
     fixed_config_path = fix_model_config(model_name)
-    
-    # Load tokenizer and model
+
+    # Load the tokenizer and model with the fixed configuration
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
         config=fixed_config_path,
-        device_map="auto"  # Use GPU if available
+        torch_dtype=torch.float16,  # Use mixed precision for faster inference
+        device_map="auto"          # Automatically allocate to GPU if available
     )
-    
+
     return pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 pipe = load_pipeline()