Update app.py

app.py

@@ -15,11 +15,16 @@ import re
 from werkzeug.utils import secure_filename
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from huggingface_hub import login
 
 app = Flask(__name__)
 
 PORT = int(os.environ.get("PORT", 7860))
 
+hf_token = os.environ.get("HF_TOKEN")
+if hf_token:
+    login(HF_TOKEN)
+
 UPLOAD_FOLDER = '/tmp/uploads'  # Change to tmp directory for Spaces
 ALLOWED_EXTENSIONS = {'py'}
 app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
@@ -28,33 +33,38 @@ os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 DATABASE_PATH = '/tmp/chat_database.db'
 
 # Initialize LangChain with Ollama LLM
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
-
-# Load model and tokenizer
-model_name = "mistralai/Mistral-7B-Instruct-v0.1"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype=torch.float16,
-    device_map="auto",
-    load_in_8bit=True
-)
-
-# Create pipeline
-pipe = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_new_tokens=512,
-    temperature=0.7,
-    top_p=0.95,
-    repetition_penalty=1.15
-)
+if hf_token:
+    model_name = "mistralai/Mistral-7B-Instruct-v0.1"
+else:
+    # Fallback to a free, smaller model
+    model_name = "microsoft/phi-4"
+
+try:
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype=torch.float16,
+        device_map="auto",
+        load_in_8bit=True
+    )
 
+    # Create pipeline
+    pipe = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        max_new_tokens=512,
+        temperature=0.7,
+        top_p=0.95,
+        repetition_penalty=1.15
+    )
 
+    # Initialize LangChain with HuggingFacePipeline
+    llm = HuggingFacePipeline(pipeline=pipe)
 
-llm = HuggingFacePipeline(pipeline=pipe)
+except Exception as e:
+    print(f"Error loading model: {e}")
+    raise
 
 @contextmanager
 def get_db_connection():