Update app.py

app.py

@@ -64,44 +64,72 @@ def get_model_name():
 def initialize_model():
     """Initialize the model with appropriate settings"""
     try:
-        model_name = get_model_name()
+        # Using a stable, free model that's known to work well in Spaces
+        model_name = "OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5"
         print(f"Initializing model: {model_name}")
 
         # Initialize tokenizer with explicit cache directory
         tokenizer = AutoTokenizer.from_pretrained(
             model_name,
             cache_dir='/tmp/transformers_cache',
-            token=os.environ.get('HUGGING_FACE_HUB_TOKEN')
+            trust_remote_code=True
         )
 
         # Initialize model with explicit cache directory
         model = AutoModelForCausalLM.from_pretrained(
             model_name,
             cache_dir='/tmp/transformers_cache',
-            token=os.environ.get('HUGGING_FACE_HUB_TOKEN'),
+            trust_remote_code=True,
             torch_dtype=torch.float16,
             device_map="auto",
             load_in_8bit=True
         )
 
-        # Create pipeline
+        # Create pipeline with specific parameters for this model
         pipe = pipeline(
             "text-generation",
             model=model,
             tokenizer=tokenizer,
             max_new_tokens=512,
+            do_sample=True,
             temperature=0.7,
             top_p=0.95,
-            repetition_penalty=1.15
+            repetition_penalty=1.15,
+            pad_token_id=tokenizer.eos_token_id
         )
 
         return HuggingFacePipeline(pipeline=pipe)
     except Exception as e:
         print(f"Error initializing model: {e}")
-        raise
+        # 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
 
-# Initialize LLM
+print("Starting model initialization...")
 llm = initialize_model()
+print("Model initialization complete!")
 
 @contextmanager
 def get_db_connection():