File size: 2,050 Bytes
65fba8a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
from fastapi import FastAPI, HTTPException, UploadFile, File, Form
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from model import SimpleTransformerModel, FullChatDataset, VoiceInterface, generate_response
import torch
import uvicorn
import os
from typing import Optional
app = FastAPI()
# CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["*"],
allow_headers=["*"],
)
# Initialize components
dataset = FullChatDataset()
model = SimpleTransformerModel(len(dataset.tokenizer))
voice_interface = VoiceInterface()
class ChatRequest(BaseModel):
prompt: str
max_length: int = 100
use_voice: bool = False
@app.post("/chat/")
async def chat_endpoint(
prompt: Optional[str] = Form(None),
max_length: int = Form(100),
use_voice: bool = Form(False),
audio_file: Optional[UploadFile] = File(None)
):
try:
# Handle voice input if audio file provided
if audio_file:
contents = await audio_file.read()
with open("temp_audio.wav", "wb") as f:
f.write(contents)
with sr.AudioFile("temp_audio.wav") as source:
audio = voice_interface.recognizer.record(source)
prompt = voice_interface.recognizer.recognize_google(audio)
os.remove("temp_audio.wav")
# If no prompt provided (either text or voice)
if not prompt:
raise HTTPException(status_code=400, detail="No input provided")
response = generate_response(
model,
dataset.tokenizer,
prompt,
max_length,
voice_interface if use_voice else None
)
return {"response": response}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.get("/")
async def read_root():
return {"message": "CyberFuture Running"}
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=8000) |