Update main.py

main.py

@@ -1,19 +1,15 @@
-import transformers
-import torch
+from transformers import pipeline, set_seed
 from fastapi import FastAPI
 from pydantic import BaseModel
 import uvicorn
 
 
-app = FastAPI()
-
+generator = pipeline('text-generation', model='gpt2')
+set_seed(42)
 
 
-model_id = "meta-llama/Meta-Llama-3-8B"
+app = FastAPI()
 
-pipeline = transformers.pipeline(
-    "text-generation", model=model_id, model_kwargs={"torch_dtype": torch.bfloat16}, device_map="auto"
-)
 
 
 
@@ -24,7 +20,7 @@ class Item(BaseModel):
 
 
 def generate(item: Item):
-    pipeline(item.prompt)
+    generator(item.prompt, max_length=30, num_return_sequences=5)
    
 
 @app.post("/generate/")