change

Dockerfile

@@ -1,10 +1,28 @@
-FROM python:3.10
+# Use the official Python image as a base
+FROM python:3.9-slim
 
+# Set environment variables
+ENV PYTHONUNBUFFERED=1 \
+    PYTHONDONTWRITEBYTECODE=1
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    git \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
+
+# Create and set the working directory
 WORKDIR /app
 
-COPY requirements.txt .
-RUN pip install -r requirements.txt
+# Copy requirements and install them
+COPY app/requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy the FastAPI application code
+COPY app /app
 
-COPY . .
+# Expose the port FastAPI will run on
+EXPOSE 8000
 
-CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]
+# Command to run the application
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]

main.py

@@ -1,20 +1,36 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
 app = FastAPI()
 
-model_name = "unsloth/Qwen2.5-7B-bnb-4bit"  # Change to your model
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+class TextGenerationRequest(BaseModel):
+    prompt: str
+    max_length: int = 100
+    temperature: float = 0.7
 
-@app.get("/")
-def read_root():
-    return {"message": "Chat API is running!"}
+# Load model and tokenizer (force CPU usage)
+model_name = "unsloth/Qwen2.5-7B-bnb-4bit"
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    trust_remote_code=True,
+    torch_dtype=torch.float32,  # Change to float32 for CPU
+    device_map="cpu"  # Force CPU usage
+)
 
-@app.post("/chat")
-def chat(prompt: str):
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=100)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return {"response": response}
+@app.post("/generate")
+async def generate_text(request: TextGenerationRequest):
+    try:
+        inputs = tokenizer(request.prompt, return_tensors="pt").to("cpu")  # Move to CPU
+        outputs = model.generate(
+            inputs.input_ids,
+            max_length=request.max_length,
+            temperature=request.temperature,
+            do_sample=True,
+        )
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return {"generated_text": generated_text}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))

requirements.txt

@@ -1,4 +1,5 @@
 fastapi
 uvicorn
-transformers
 torch
+transformers
+accelerate