Spaces:
Paused
Paused
| from fastapi import FastAPI, HTTPException | |
| from pydantic import BaseModel | |
| from langchain_huggingface import ChatHuggingFace, HuggingFacePipeline | |
| import uvicorn | |
| import os | |
| # Define request model | |
| class PromptRequest(BaseModel): | |
| prompt: str | |
| app = FastAPI(title="Local LLM Server") | |
| # Global variable for model | |
| model = None | |
| async def startup_event(): | |
| global model | |
| print("Initializing TinyLlama... This will auto-download the model if not present.") | |
| # Initialize the pipeline as requested | |
| llm = HuggingFacePipeline.from_model_id( | |
| model_id="TinyLlama/TinyLlama-1.1B-Chat-v1.0", | |
| task="text-generation", | |
| pipeline_kwargs=dict( | |
| temperature=0.7, | |
| max_new_tokens=512, | |
| do_sample=True | |
| ) | |
| ) | |
| model = ChatHuggingFace(llm=llm) | |
| print("Model loaded and ready to serve!") | |
| async def generate_text(request: PromptRequest): | |
| if not model: | |
| raise HTTPException(status_code=503, detail="Model is loading, please wait...") | |
| try: | |
| # Invoke the model | |
| result = model.invoke(request.prompt) | |
| # Return the content of the AIMessage | |
| return {"response": result.content} | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| if __name__ == "__main__": | |
| print("Starting server on http://127.0.0.1:8000") | |
| uvicorn.run(app, host="127.0.0.1", port=8000) | |