Setup llm

main.py

@@ -1,7 +1,20 @@
 from fastapi import FastAPI
+from pydantic import BaseModel
+from transformers import pipeline
+
+# Load a small model to test (can change to bigger one later)
+generator = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.2")
 
 app = FastAPI()
 
+class Prompt(BaseModel):
+    text: str
+
 @app.get("/")
 def read_root():
     return {"Hello": "docker embeddings!"}
+
+@app.post("/generate")
+async def generate_text(prompt: Prompt):
+    output = generator(prompt.text, max_new_tokens=200, do_sample=True, temperature=0.7)
+    return {"response": output[0]["generated_text"]}

requirements.txt

@@ -4,3 +4,6 @@ sentencepiece==0.1.*
 torch==1.11.*
 transformers==4.*
 uvicorn[standard]==0.17.*
+transformers
+torch
+accelerate