updated with a smaller model.

app/main.py

@@ -1,11 +1,14 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from vllm import LLM, SamplingParams
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
 app = FastAPI()
 
 # Initialize the model (we'll use a small model for this example)
-model = LLM(model="EleutherAI/gpt-neo-125M")
+model_name = "EleutherAI/gpt-neo-125M"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
 class GenerateRequest(BaseModel):
     prompt: str
@@ -13,12 +16,13 @@ class GenerateRequest(BaseModel):
 @app.post("/generate")
 async def generate(request: GenerateRequest):
     try:
-        sampling_params = SamplingParams(temperature=0.7, max_tokens=100)
-        outputs = model.generate([request.prompt], sampling_params)
-        return {"generated_text": outputs[0].outputs[0].text}
+        input_ids = tokenizer.encode(request.prompt, return_tensors="pt")
+        output = model.generate(input_ids, max_length=100, num_return_sequences=1)
+        generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+        return {"generated_text": generated_text}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.get("/")
 async def root():
-    return {"message": "vLLM server is running"}
+    return {"message": "Model server is running"}