Update app.py

app.py

@@ -2,6 +2,7 @@ from flask import Flask, render_template, request, jsonify
 import subprocess
 import tempfile
 import os
+import shutil
 from langchain_community.llms import HuggingFacePipeline
 from langchain.prompts import PromptTemplate
 from langchain.chains import LLMChain
@@ -15,167 +16,118 @@ import re
 from werkzeug.utils import secure_filename
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from huggingface_hub import login
 
 app = Flask(__name__)
 
 # Configuration for Hugging Face Spaces
 PORT = int(os.environ.get("PORT", 7860))
 
-# Set cache directories to /tmp
-os.environ['TRANSFORMERS_CACHE'] = '/tmp/transformers_cache'
-os.environ['HF_HOME'] = '/tmp/hf_home'
-os.environ['XDG_CACHE_HOME'] = '/tmp/cache'
-os.environ['HF_DATASETS_CACHE'] = '/tmp/datasets_cache'
-
-# Create necessary directories with proper permissions
-for directory in [
-    '/tmp/transformers_cache',
-    '/tmp/hf_home',
-    '/tmp/cache',
-    '/tmp/datasets_cache',
-    '/tmp/uploads'
-]:
-    os.makedirs(directory, exist_ok=True)
-
-# Configure upload folder inside the space
+# Create and set up a writable directory in /tmp
+CACHE_DIR = "/tmp/huggingface_cache"
+os.makedirs(CACHE_DIR, exist_ok=True)
+os.environ['TRANSFORMERS_CACHE'] = CACHE_DIR
+os.environ['HF_HOME'] = CACHE_DIR
+os.environ['XDG_CACHE_HOME'] = CACHE_DIR
+
+# Configure upload folder
 UPLOAD_FOLDER = '/tmp/uploads'
 ALLOWED_EXTENSIONS = {'py'}
 app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
+os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 
 # Database configuration
 DATABASE_PATH = '/tmp/chat_database.db'
 
-def get_model_name():
-    """Determine which model to use based on token availability"""
-    try:
-        hf_token = os.environ.get("HF_TOKEN")
-        if hf_token:
-            # Set token in environment and return gated model name
-            os.environ['HUGGING_FACE_HUB_TOKEN'] = hf_token
-            return "mistralai/Mistral-7B-Instruct-v0.1"
-        else:
-            # Return free model if no token
-            return "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-    except Exception as e:
-        print(f"Error accessing token: {e}")
-        return "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-
 def initialize_model():
     """Initialize the model with appropriate settings"""
     try:
-        # Using a stable, free model that's known to work well in Spaces
-        model_name = "OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5"
+        # Use a smaller model that's more likely to work in the Space
+        model_name = "facebook/opt-350m"
         print(f"Initializing model: {model_name}")
 
-        # Initialize tokenizer with explicit cache directory
+        # Initialize tokenizer
         tokenizer = AutoTokenizer.from_pretrained(
             model_name,
-            cache_dir='/tmp/transformers_cache',
-            trust_remote_code=True
+            cache_dir=CACHE_DIR,
+            local_files_only=False
         )
 
-        # Initialize model with explicit cache directory
+        # Initialize model with minimal settings
         model = AutoModelForCausalLM.from_pretrained(
             model_name,
-            cache_dir='/tmp/transformers_cache',
-            trust_remote_code=True,
+            cache_dir=CACHE_DIR,
+            local_files_only=False,
             torch_dtype=torch.float16,
-            device_map="auto",
-            load_in_8bit=True
+            low_cpu_mem_usage=True
         )
 
-        # Create pipeline with specific parameters for this model
+        # Create pipeline
         pipe = pipeline(
             "text-generation",
             model=model,
             tokenizer=tokenizer,
-            max_new_tokens=512,
-            do_sample=True,
+            max_new_tokens=256,
             temperature=0.7,
             top_p=0.95,
             repetition_penalty=1.15,
-            pad_token_id=tokenizer.eos_token_id
+            device_map="auto"
         )
 
         return HuggingFacePipeline(pipeline=pipe)
     except Exception as e:
         print(f"Error initializing model: {e}")
-        # If the main model fails, try an even smaller fallback
-        try:
-            model_name = "facebook/opt-125m"
-            print(f"Trying fallback model: {model_name}")
-            
-            tokenizer = AutoTokenizer.from_pretrained(model_name)
-            model = AutoModelForCausalLM.from_pretrained(
-                model_name,
-                torch_dtype=torch.float16,
-                device_map="auto"
-            )
-            
-            pipe = pipeline(
-                "text-generation",
-                model=model,
-                tokenizer=tokenizer,
-                max_new_tokens=512,
-                temperature=0.7,
-                top_p=0.95
-            )
-            
-            return HuggingFacePipeline(pipeline=pipe)
-        except Exception as fallback_error:
-            print(f"Fallback model also failed: {fallback_error}")
-            raise
-
-print("Starting model initialization...")
-llm = initialize_model()
-print("Model initialization complete!")
-
-@contextmanager
-def get_db_connection():
-    conn = sqlite3.connect(DATABASE_PATH)
-    conn.row_factory = sqlite3.Row
-    try:
-        yield conn
-    finally:
-        conn.close()
-
+        raise
 
+# Initialize database
 def init_db():
-    with get_db_connection() as conn:
-        conn.execute('''
-            CREATE TABLE IF NOT EXISTS chats (
-                id TEXT PRIMARY KEY,
-                title TEXT,
-                date TEXT,
-                last_message TEXT
-            )
-        ''')
-
-        conn.execute('''
-            CREATE TABLE IF NOT EXISTS messages (
-                id INTEGER PRIMARY KEY AUTOINCREMENT,
-                chat_id TEXT,
-                role TEXT,
-                content TEXT,
-                timestamp TEXT,
-                FOREIGN KEY (chat_id) REFERENCES chats (id)
-            )
-        ''')
-
-        conn.execute('''
-            CREATE TABLE IF NOT EXISTS important_info (
-                id INTEGER PRIMARY KEY AUTOINCREMENT,
-                chat_id TEXT,
-                content TEXT,
-                FOREIGN KEY (chat_id) REFERENCES chats (id)
-            )
-        ''')
-        conn.commit()
-
+    """Initialize the database"""
+    try:
+        with get_db_connection() as conn:
+            conn.execute('''
+                CREATE TABLE IF NOT EXISTS chats (
+                    id TEXT PRIMARY KEY,
+                    title TEXT,
+                    date TEXT,
+                    last_message TEXT
+                )
+            ''')
+
+            conn.execute('''
+                CREATE TABLE IF NOT EXISTS messages (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    chat_id TEXT,
+                    role TEXT,
+                    content TEXT,
+                    timestamp TEXT,
+                    FOREIGN KEY (chat_id) REFERENCES chats (id)
+                )
+            ''')
+
+            conn.execute('''
+                CREATE TABLE IF NOT EXISTS important_info (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    chat_id TEXT,
+                    content TEXT,
+                    FOREIGN KEY (chat_id) REFERENCES chats (id)
+                )
+            ''')
+            conn.commit()
+    except Exception as e:
+        print(f"Error initializing database: {e}")
+        raise
 
-# Initialize database on startup
-init_db()
+# Initialize the application
+try:
+    print("Initializing database...")
+    init_db()
+    print("Database initialized successfully")
+    
+    print("Starting model initialization...")
+    llm = initialize_model()
+    print("Model initialized successfully")
+except Exception as e:
+    print(f"Initialization error: {e}")
+    raise
 
 
 class ChatSession: