Update app.py

app.py

@@ -1,13 +1,12 @@
+import os
+import torch
 from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.middleware.cors import CORSMiddleware
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from autoawq import AutoAWQForCausalLM  # Add this import
 from transformers import pipeline
-import torch
-import os
 import time
 
 # Ensure cache directories exist
@@ -17,26 +16,28 @@ os.makedirs(os.getenv('MPLCONFIGDIR', '/app/cache/matplotlib'), exist_ok=True)
 # Initialize FastAPI app
 app = FastAPI()
 
-# Mount Gradio app
-gradio_app = gr.Blocks()
-
-# Model loading function
+# Model loading function without autoawq
 def load_model():
     model_name = "trillionlabs/Trillion-7B-preview-AWQ"
     
-    # Load tokenizer with special handling
+    # Load tokenizer with error handling
     try:
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-    except:
-        # Fallback to using the model's tokenizer.json directly
-        from transformers import PreTrainedTokenizerFast
-        tokenizer = PreTrainedTokenizerFast(tokenizer_file=f"{model_name}/tokenizer.json")
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_name,
+            trust_remote_code=True
+        )
+    except Exception as e:
+        print(f"Error loading tokenizer: {e}")
+        # Fallback to a more basic tokenizer if needed
+        from transformers import LlamaTokenizer
+        tokenizer = LlamaTokenizer.from_pretrained(model_name)
     
-    # Load model with CPU support
+    # Load model with CPU configuration
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
         device_map="cpu",
-        torch_dtype=torch.float32
+        torch_dtype=torch.float32,
+        trust_remote_code=True
     )
     
     # Create text generation pipeline
@@ -49,8 +50,13 @@ def load_model():
     
     return text_generator
 
-# Load model (this will happen when the server starts)
-text_generator = load_model()
+# Load model
+try:
+    text_generator = load_model()
+except Exception as e:
+    print(f"Failed to load model: {e}")
+    # You might want to exit here or load a smaller model instead
+    raise
 
 # API endpoint for text generation
 @app.post("/api/generate")
@@ -58,9 +64,9 @@ async def generate_text(request: Request):
     try:
         data = await request.json()
         prompt = data.get("prompt", "")
-        max_length = data.get("max_length", 100)
+        max_length = min(int(data.get("max_length", 100)), 500)  # Limit to 500 tokens
         
-        # Generate text
+        # Generate text with timing
         start_time = time.time()
         outputs = text_generator(
             prompt,
@@ -68,14 +74,15 @@ async def generate_text(request: Request):
             do_sample=True,
             temperature=0.7,
             top_k=50,
-            top_p=0.95
+            top_p=0.95,
+            pad_token_id=0  # Might be needed for some models
         )
         generation_time = time.time() - start_time
         
         return JSONResponse({
             "generated_text": outputs[0]["generated_text"],
-            "generation_time": generation_time,
-            "model": "trillionlabs/Trillion-7B-preview-AWQ",
+            "generation_time": round(generation_time, 2),
+            "model": "Trillion-7B-preview-AWQ",
             "device": "cpu"
         })
     except Exception as e:
@@ -83,26 +90,58 @@ async def generate_text(request: Request):
 
 # Gradio interface
 def gradio_generate(prompt, max_length=100):
-    outputs = text_generator(
-        prompt,
-        max_length=max_length,
-        do_sample=True,
-        temperature=0.7,
-        top_k=50,
-        top_p=0.95
-    )
-    return outputs[0]["generated_text"]
+    try:
+        max_length = min(int(max_length), 500)  # Limit to 500 tokens
+        outputs = text_generator(
+            prompt,
+            max_length=max_length,
+            do_sample=True,
+            temperature=0.7,
+            top_k=50,
+            top_p=0.95,
+            pad_token_id=0
+        )
+        return outputs[0]["generated_text"]
+    except Exception as e:
+        return f"Error generating text: {str(e)}"
 
-with gradio_app:
-    gr.Markdown("# Trillion-7B-preview-AWQ Demo (CPU)")
-    gr.Markdown("This is a CPU-only demo of the Trillion-7B-preview-AWQ model running with 16GB RAM.")
+with gr.Blocks() as gradio_app:
+    gr.Markdown("""
+    # Trillion-7B-preview-AWQ Demo (CPU)
+    *Running on CPU with 16GB RAM - responses may be slow*
+    """)
+    
+    with gr.Row():
+        input_prompt = gr.Textbox(
+            label="Input Prompt", 
+            lines=5,
+            placeholder="Enter your prompt here..."
+        )
+        output_text = gr.Textbox(
+            label="Generated Text", 
+            lines=5,
+            interactive=False
+        )
     
     with gr.Row():
-        input_prompt = gr.Textbox(label="Input Prompt", lines=5)
-        output_text = gr.Textbox(label="Generated Text", lines=5)
+        length_slider = gr.Slider(
+            minimum=50,
+            maximum=500,
+            value=100,
+            step=10,
+            label="Max Length"
+        )
+        generate_btn = gr.Button("Generate", variant="primary")
     
-    length_slider = gr.Slider(50, 500, value=100, label="Max Length")
-    generate_btn = gr.Button("Generate")
+    # Additional examples
+    examples = gr.Examples(
+        examples=[
+            ["Explain quantum computing in simple terms."],
+            ["Write a short poem about artificial intelligence."],
+            ["How do I make a good cup of coffee?"]
+        ],
+        inputs=input_prompt
+    )
     
     generate_btn.click(
         fn=gradio_generate,
@@ -120,4 +159,9 @@ app.add_middleware(
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
-)
+)
+
+# Health check endpoint
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy", "model_loaded": text_generator is not None}