Update app.py

app.py

@@ -1,12 +1,12 @@
 import os
 import torch
+import time
+import warnings
 from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-import time
-import warnings
 
 # Suppress specific warnings
 warnings.filterwarnings("ignore", category=FutureWarning, module="transformers.utils.hub")
@@ -44,24 +44,12 @@ def load_model():
     
     log_message("Loading model...")
     try:
-        # Try loading with IPEX optimization if available
-        try:
-            import intel_extension_for_pytorch as ipex
-            use_ipex = True
-        except ImportError:
-            use_ipex = False
-            log_message("IPEX not available, using standard CPU version")
-
         model = AutoModelForCausalLM.from_pretrained(
             model_name,
             torch_dtype=torch.float32,
             trust_remote_code=True
         )
         
-        if use_ipex:
-            log_message("Applying IPEX optimization...")
-            model = ipex.optimize(model)
-        
         # Explicitly move to CPU
         model = model.to("cpu")
         model.eval()
@@ -78,12 +66,12 @@ def load_model():
         device="cpu"
     )
     
-    return text_generator
+    return text_generator, tokenizer
 
 # Load model
 try:
     log_message("Starting model loading process...")
-    text_generator = load_model()
+    text_generator, tokenizer = load_model()
     log_message("Model loaded successfully")
 except Exception as e:
     log_message(f"Critical error loading model: {e}")
@@ -109,17 +97,18 @@ async def api_generate(request: Request):
             temperature=0.7,
             top_k=50,
             top_p=0.95,
-            pad_token_id=tokenizer.eos_token_id if hasattr(text_generator, 'tokenizer') else 0
+            pad_token_id=tokenizer.eos_token_id
         )
         generation_time = time.time() - start_time
         
-        return JSONResponse({
+        response_data = {
             "generated_text": outputs[0]["generated_text"],
             "time_seconds": round(generation_time, 2),
-            "tokens_generated": len(text_generator.tokenizer.tokenize(outputs[0]["generated_text"]) if hasattr(text_generator, 'tokenizer') else None,
+            "tokens_generated": len(tokenizer.tokenize(outputs[0]["generated_text"])),
             "model": "Trillion-7B-preview-AWQ",
             "device": "cpu"
-        })
+        }
+        return JSONResponse(response_data)
     except Exception as e:
         log_message(f"API Error: {e}")
         return JSONResponse({"error": str(e)}, status_code=500)
@@ -149,7 +138,7 @@ def gradio_generate(prompt, max_length=100):
             temperature=0.7,
             top_k=50,
             top_p=0.95,
-            pad_token_id=tokenizer.eos_token_id if hasattr(text_generator, 'tokenizer') else 0
+            pad_token_id=tokenizer.eos_token_id
         )
         return outputs[0]["generated_text"]
     except Exception as e: