Update app.py

app.py

@@ -1,20 +1,20 @@
 import streamlit as st
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
-import torch
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, LlamaConfig
 import json
 import os
+import requests
+import torch
 
 # Streamlit app configuration
 st.set_page_config(page_title="AI Chatbot", layout="centered")
 
-# Fix the model's configuration dynamically
+# Fix and modify the model configuration dynamically
 def fix_model_config(model_name):
-    # Load the configuration file directly from the Hugging Face hub
+    # Download the configuration file 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()  # Ensure the request is successful
         with open(config_path, "w") as f:
@@ -27,7 +27,7 @@ def fix_model_config(model_name):
     # Fix the `rope_scaling` field
     if "rope_scaling" in config:
         config["rope_scaling"] = {
-            "type": "linear",  # Replace the problematic structure with supported format
+            "type": "linear",  # Only keep 'type' and 'factor'
             "factor": config["rope_scaling"].get("factor", 1.0)
         }
 
@@ -45,12 +45,12 @@ def load_pipeline():
     # Fix the model configuration
     fixed_config_path = fix_model_config(model_name)
 
-    # Load the tokenizer and model with the fixed configuration
+    # Use the fixed configuration to load the model
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
         config=fixed_config_path,
-        torch_dtype=torch.float16,  # Use mixed precision for faster inference
+        torch_dtype=torch.float16,  # Mixed precision for efficiency
         device_map="auto"          # Automatically allocate to GPU if available
     )