mirror of
https://github.com/theroyallab/tabbyAPI.git
synced 2026-03-14 15:57:27 +00:00
40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
# main.py
|
|
|
|
from fastapi import FastAPI, HTTPException
|
|
from pydantic import BaseModel
|
|
from llm import ModelManager
|
|
from uvicorn import run
|
|
|
|
app = FastAPI()
|
|
|
|
# Example: Using a different model directory
|
|
modelManager = ModelManager("/home/david/Models/SynthIA-7B-v2.0-5.0bpw-h6-exl2")
|
|
|
|
class TextRequest(BaseModel):
|
|
model: str
|
|
messages: list[dict]
|
|
temperature: float
|
|
|
|
class TextResponse(BaseModel):
|
|
response: str
|
|
generation_time: str
|
|
|
|
@app.post("/generate-text", response_model=TextResponse)
|
|
def generate_text(request: TextRequest):
|
|
try:
|
|
#model_path = request.model # You can use this path to load a specific model if needed
|
|
messages = request.messages
|
|
#temperature = request.temperature
|
|
|
|
# Assuming you need to extract the user's message from the messages list
|
|
user_message = next(msg["content"] for msg in messages if msg["role"] == "user")
|
|
|
|
# You can then use user_message as the prompt for generation
|
|
output, generation_time = modelManager.generate_text(user_message)
|
|
return {"response": output, "generation_time": generation_time}
|
|
except RuntimeError as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
if __name__ == "__main__":
|
|
run(app, host="0.0.0.0", port=8012, reload=True)
|