Create app.py

app.py

@@ -0,0 +1,113 @@
+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 transformers import pipeline
+import torch
+import os
+import time
+
+# Initialize FastAPI app
+app = FastAPI()
+
+# Mount Gradio app
+gradio_app = gr.Blocks()
+
+# Model loading function
+def load_model():
+    model_name = "trillionlabs/Trillion-7B-preview-AWQ"
+    
+    # Load tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    
+    # Load model with CPU support
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        device_map="cpu",
+        torch_dtype=torch.float32
+    )
+    
+    # Create text generation pipeline
+    text_generator = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        device="cpu"
+    )
+    
+    return text_generator
+
+# Load model (this will happen when the server starts)
+text_generator = load_model()
+
+# API endpoint for text generation
+@app.post("/api/generate")
+async def generate_text(request: Request):
+    try:
+        data = await request.json()
+        prompt = data.get("prompt", "")
+        max_length = data.get("max_length", 100)
+        
+        # Generate text
+        start_time = time.time()
+        outputs = text_generator(
+            prompt,
+            max_length=max_length,
+            do_sample=True,
+            temperature=0.7,
+            top_k=50,
+            top_p=0.95
+        )
+        generation_time = time.time() - start_time
+        
+        return JSONResponse({
+            "generated_text": outputs[0]["generated_text"],
+            "generation_time": generation_time,
+            "model": "trillionlabs/Trillion-7B-preview-AWQ",
+            "device": "cpu"
+        })
+    except Exception as e:
+        return JSONResponse({"error": str(e)}, status_code=500)
+
+# 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"]
+
+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.Row():
+        input_prompt = gr.Textbox(label="Input Prompt", lines=5)
+        output_text = gr.Textbox(label="Generated Text", lines=5)
+    
+    length_slider = gr.Slider(50, 500, value=100, label="Max Length")
+    generate_btn = gr.Button("Generate")
+    
+    generate_btn.click(
+        fn=gradio_generate,
+        inputs=[input_prompt, length_slider],
+        outputs=output_text
+    )
+
+# Mount Gradio app
+app = gr.mount_gradio_app(app, gradio_app, path="/")
+
+# CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)