added server

Dockerfile

@@ -0,0 +1,18 @@
+# Use an official Python image as a base
+FROM python:3.12-slim
+
+# Set the working directory inside the container
+WORKDIR /app
+
+# Copy the requirements file and install dependencies
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy the rest of the application code
+COPY . .
+
+# Expose port 8000 for the FastAPI app
+EXPOSE 8000
+
+# Run the FastAPI app with uvicorn
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]

main.py

@@ -0,0 +1,52 @@
+import os
+import torch
+from fastapi import FastAPI, Request
+from pydantic import BaseModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from dotenv import load_dotenv
+
+# Load environment variables from a .env file (useful for local development)
+load_dotenv()
+
+# Initialize FastAPI app
+app = FastAPI()
+
+# Set your Hugging Face token from environment variable
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+MODEL = "meta-llama/Llama-3.2-3B-Instruct"
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
+
+# Load model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained(MODEL, token=HF_TOKEN)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL,
+    token=HF_TOKEN,
+    torch_dtype=torch.float16,  # Use float16 for better GPU memory usage
+    device_map="auto"
+)
+
+# Pydantic model for input
+class PromptRequest(BaseModel):
+    prompt: str
+    max_new_tokens: int = 100
+    temperature: float = 0.7
+
+@app.post("/generate/")
+async def generate_text(request: PromptRequest):
+    inputs = tokenizer(request.prompt, return_tensors="pt").to(device)
+    
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=request.max_new_tokens,
+            temperature=request.temperature,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
+    
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return {"response": response}
+

requirements.txt

@@ -0,0 +1,5 @@
+fastapi
+transformers
+torch
+uvicorn
+python-dotenv