Update app.py

app.py

@@ -1,45 +1,35 @@
 from flask import Flask, render_template, request, jsonify
-import subprocess
-import tempfile
 import os
-import shutil
-from langchain_community.llms import HuggingFacePipeline
+import tempfile
+from pathlib import Path
+from langchain_community.llms import HuggingFaceHub
 from langchain.prompts import PromptTemplate
 from langchain.chains import LLMChain
 from langchain.memory import ConversationBufferMemory
 from datetime import datetime
-import json
-from typing import Dict, List
 import sqlite3
 from contextlib import contextmanager
-import re
 from werkzeug.utils import secure_filename
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
+# Initialize Flask application
 app = Flask(__name__)
 
-# Configuration for Hugging Face Spaces
+# Configuration
 PORT = int(os.environ.get("PORT", 7860))
-
-# 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
+UPLOAD_DIR = "/tmp/uploads"
+DATABASE_PATH = "/tmp/chat_database.db"
 
-# Configure upload folder
-UPLOAD_FOLDER = '/tmp/uploads'
-ALLOWED_EXTENSIONS = {'py'}
-app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
-os.makedirs(UPLOAD_FOLDER, exist_ok=True)
+# Create necessary directories
+for dir_path in [CACHE_DIR, UPLOAD_DIR]:
+    Path(dir_path).mkdir(parents=True, exist_ok=True)
 
-# Database configuration
-DATABASE_PATH = '/tmp/chat_database.db'
+# Set environment variables
+os.environ['TRANSFORMERS_CACHE'] = CACHE_DIR
+os.environ['HUGGINGFACE_HUB_CACHE'] = CACHE_DIR
 
-llm = None
-llm_chain = None
+# Configure Flask app
+app.config['UPLOAD_FOLDER'] = UPLOAD_DIR
 
 # Initialize prompt template
 prompt_template = """
@@ -59,15 +49,14 @@ Output Guidelines:
    - Use ```python for code blocks
    - Use `code` for inline code references
    - Provide raw text without HTML formatting
-   - Strictly include explanation only after code blocks
+   - Include explanations after code blocks
 
 2. Code Organization:
-   - Default to single, focused code snippets for clarity
-   - Only split into multiple snippets(each individually runnable) if:
-     a) Multiple distinct concepts are requested
-     b) Complex functionality requires modular explanation
-     
-   - Mark critical information with [IMPORTANT] prefix and give small explanations with some bold headings if required and in white font always.
+   - Default to single, focused code snippets
+   - Only split into multiple snippets if necessary
+   - Mark critical information with [IMPORTANT] prefix
+
+Please provide a clear and detailed response:
 """
 
 prompt = PromptTemplate(
@@ -85,47 +74,59 @@ def get_db_connection():
     finally:
         conn.close()
 
-def initialize_model():
-    """Initialize the model with appropriate settings"""
-    try:
-        # 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
-        tokenizer = AutoTokenizer.from_pretrained(
-            model_name,
-            cache_dir=CACHE_DIR,
-            local_files_only=False
-        )
-
-        # Initialize model with minimal settings
-        model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            cache_dir=CACHE_DIR,
-            local_files_only=False,
-            torch_dtype=torch.float16,
-            low_cpu_mem_usage=True
-        )
+def init_db():
+    """Initialize the database"""
+    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()
 
-        # Create pipeline
-        pipe = pipeline(
-            "text-generation",
-            model=model,
-            tokenizer=tokenizer,
-            max_new_tokens=256,
-            temperature=0.7,
-            top_p=0.95,
-            repetition_penalty=1.15,
-            device_map="auto"
+def initialize_llm():
+    """Initialize the language model"""
+    try:
+        # Using Mistral model through Hugging Face Hub
+        llm = HuggingFaceHub(
+            repo_id="mistralai/Mistral-7B-Instruct-v0.1",
+            huggingfacehub_api_token=os.environ.get('HUGGINGFACE_API_TOKEN'),
+            model_kwargs={
+                "temperature": 0.7,
+                "max_new_tokens": 512,
+                "top_p": 0.95,
+                "repetition_penalty": 1.15
+            }
         )
-
-        return HuggingFacePipeline(pipeline=pipe)
+        return llm
     except Exception as e:
-        print(f"Error initializing model: {e}")
+        print(f"Error initializing LLM: {e}")
         raise
 
-
+        
 class ChatSession:
     def __init__(self, session_id):
         self.session_id = session_id
@@ -230,65 +231,28 @@ class ChatSession:
 
         self.important_info = []
 
+print("Starting application initialization...")
 
-def init_db():
-    """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
+try:
+    print("Initializing database...")
+    init_db()
+    print("Database initialized successfully")
 
-def initialize_app():
-    """Initialize all components of the application"""
-    global llm, llm_chain
-    
-    try:
-        print("Initializing database...")
-        init_db()
-        print("Database initialized successfully")
-        
-        print("Starting model initialization...")
-        llm = initialize_model()
-        print("Model initialized successfully")
-        
-        print("Creating LLM chain...")
-        llm_chain = LLMChain(llm=llm, prompt=prompt)
-        print("LLM chain created successfully")
-    except Exception as e:
-        print(f"Initialization error: {e}")
-        raise
+    print("Initializing LLM...")
+    llm = initialize_llm()
+    if llm is None:
+        raise ValueError("LLM initialization failed")
+    print("LLM initialized successfully")
 
-        
+    print("Creating LLM chain...")
+    llm_chain = LLMChain(llm=llm, prompt=prompt)
+    print("LLM chain created successfully")
+
+except Exception as e:
+    print(f"Fatal initialization error: {e}")
+    raise
+
+    
 # Rest of the prompt template and other configurations remain the same
 prompt_template = """
 Role: You are Figr Code Assistant, specializing in providing clear, error-free Python code solutions.
@@ -468,11 +432,10 @@ def get_chat_list():
 
 @app.route("/api/chat", methods=["POST"])
 def chat():
-    global llm_chain
-    if llm_chain is None:
+    if not llm_chain:
         return jsonify({
-            "response": "System is still initializing. Please try again in a moment.",
-            "success": False
+            "success": False,
+            "response": "System is still initializing. Please try again in a moment."
         })
     data = request.json
     user_input = data.get("message", "")
@@ -720,7 +683,7 @@ def home():
     return render_template("index.html")
 
 
-initialize_app()
+
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=PORT)