Spaces:
Runtime error
Runtime error
| import os | |
| from fastapi import FastAPI, Header, HTTPException | |
| from pydantic import BaseModel | |
| from transformers import pipeline | |
| API_KEY = os.getenv("API_KEY") | |
| app = FastAPI() | |
| generator = pipeline( | |
| "text-generation", | |
| model="mistralai/Mistral-7B-Instruct-v0.3", | |
| max_new_tokens=256 | |
| ) | |
| class PromptRequest(BaseModel): | |
| prompt: str | |
| def generate(data: PromptRequest, x_api_key: str = Header(None)): | |
| if API_KEY is None: | |
| raise HTTPException(status_code=500, detail="API key not configured") | |
| if x_api_key != API_KEY: | |
| raise HTTPException(status_code=401, detail="Invalid API Key") | |
| result = generator(data.prompt) | |
| return {"response": result[0]["generated_text"]} | |