Update app.py

app.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, LlamaConfig
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 import json
 import os
 import requests
@@ -10,32 +10,27 @@ st.set_page_config(page_title="AI Chatbot", layout="centered")
 
 # Fix and modify the model configuration dynamically
 def fix_model_config(model_name):
-    # 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"
+    fixed_config_path = "fixed_config.json"
 
-    if not os.path.exists(config_path):
+    # Download and modify config.json
+    if not os.path.exists(fixed_config_path):
         response = requests.get(config_url)
-        response.raise_for_status()  # Ensure the request is successful
-        with open(config_path, "w") as f:
-            f.write(response.text)
+        response.raise_for_status()
+        config = response.json()
 
-    # 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",
+                "factor": config["rope_scaling"].get("factor", 1.0)
+            }
 
-    # Fix the `rope_scaling` field
-    if "rope_scaling" in config:
-        config["rope_scaling"] = {
-            "type": "linear",  # Only keep 'type' and 'factor'
-            "factor": config["rope_scaling"].get("factor", 1.0)
-        }
+        # Save the fixed config
+        with open(fixed_config_path, "w") as f:
+            json.dump(config, f)
 
-    # Save the fixed configuration locally
-    with open(config_path, "w") as f:
-        json.dump(config, f)
-
-    return config_path
+    return fixed_config_path
 
 # Load the pipeline
 @st.cache_resource
@@ -45,15 +40,16 @@ def load_pipeline():
     # Fix the model configuration
     fixed_config_path = fix_model_config(model_name)
 
-    # Use the fixed configuration to load the model
+    # Load tokenizer and model
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
         config=fixed_config_path,
-        torch_dtype=torch.float16,  # Mixed precision for efficiency
-        device_map="auto"          # Automatically allocate to GPU if available
+        torch_dtype=torch.float16,
+        device_map="auto"
     )
 
+    # Return the text generation pipeline
     return pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 pipe = load_pipeline()