update app.py

app.py

@@ -3,7 +3,7 @@ import gradio as gr
 import torch, pandas as pd
 import matplotlib.pyplot as plt
 import seaborn as sns
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
 # ZeroGPU support
 try:
@@ -24,12 +24,41 @@ except ImportError:
 MODEL_NAME = "fdtn-ai/Foundation-Sec-8B"
 #MODEL_NAME = "sshleifer/tiny-gpt2"
 
-# Initialize tokenizer and model
+# Initialize tokenizer and model using pipeline approach
 print(f"Loading model: {MODEL_NAME}")
-tok = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME, torch_dtype=torch.float16, device_map="auto"
-).eval()
+try:
+    print(f"Initializing text generation model: {MODEL_NAME}")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    text_pipeline = pipeline(
+        "text-generation",
+        model=MODEL_NAME,
+        tokenizer=tokenizer,
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+        trust_remote_code=True
+    )
+    print(f"Model initialized successfully: {MODEL_NAME}")
+    
+    # Extract model and tokenizer from pipeline for direct access
+    model = text_pipeline.model
+    tok = text_pipeline.tokenizer
+    
+except Exception as e:
+    print(f"Error initializing model {MODEL_NAME}: {str(e)}")
+    print("Falling back to tiny-gpt2...")
+    MODEL_NAME = "sshleifer/tiny-gpt2"
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    text_pipeline = pipeline(
+        "text-generation",
+        model=MODEL_NAME,
+        tokenizer=tokenizer,
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+        trust_remote_code=True
+    )
+    model = text_pipeline.model
+    tok = text_pipeline.tokenizer
+    print(f"Fallback model loaded: {MODEL_NAME}")
 
 # Log device information
 if hasattr(model, 'device'):