Integrate meta-llama model with FastAPI

Dockerfile

@@ -1,6 +1,3 @@
-# Read the doc: https://huggingface.co/docs/hub/spaces-sdks-docker
-# you will also find guides on how best to write your Dockerfile
-
 FROM python:3.9
 
 RUN useradd -m -u 1000 user
@@ -9,8 +6,16 @@ ENV PATH="/home/user/.local/bin:$PATH"
 
 WORKDIR /app
 
+# Copy dependencies and install them
 COPY --chown=user ./requirements.txt requirements.txt
 RUN pip install --no-cache-dir --upgrade -r requirements.txt
 
+# Pre-load the model to save time during runtime
+RUN python -c "from transformers import AutoTokenizer, AutoModelForCausalLM; \
+    AutoTokenizer.from_pretrained('meta-llama/Llama-3.1-8B-Instruct'); \
+    AutoModelForCausalLM.from_pretrained('meta-llama/Llama-3.1-8B-Instruct')"
+
 COPY --chown=user . /app
+
+# Run the app using uvicorn
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -1,3 +1,22 @@
-import gradio as gr
+from fastapi import FastAPI
+from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-gr.load("models/meta-llama/Llama-3.1-8B-Instruct").launch()
+# Load the model and tokenizer
+model_name = "meta-llama/Llama-3.1-8B-Instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+
+# Initialize FastAPI
+app = FastAPI()
+
+# Define request and response models
+class Prompt(BaseModel):
+    text: str
+
+@app.post("/generate")
+def generate_text(prompt: Prompt):
+    inputs = tokenizer(prompt.text, return_tensors="pt")
+    outputs = model.generate(**inputs, max_length=100)
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return {"generated_text": generated_text}

requirements.txt

@@ -1,2 +1,4 @@
 fastapi
 uvicorn[standard]
+transformers
+torch