Upload 3 files

app.py

@@ -1,27 +1,31 @@
 from fastapi import FastAPI
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
 
 ## Create FastAPI app instance
 app = FastAPI()
 
 # Initialize a text generation pipeline
-# Load model directly
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-
+# Load model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
 model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
 
+# Check for GPU availability
+device = 0 if torch.cuda.is_available() else -1
+
 @app.get("/")
 def home():
     return {"message": "Welcome to the Text Generation API!"}
 
 #Define a function to handle the GET request at /generate 
 @app.get("/generate")
-def generate(text:str):
+def generate(text: str):
     ## Generate text using the model
-    output=pipe(text)
+    input_ids = tokenizer(text, return_tensors="pt").input_ids
+    outputs = model.generate(input_ids, max_length=100)
+    output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     
-    ##return the text generate in Json response
-    return {"output":output[0]['generated_text']}
+    ##return the text generated in Json response
+    return {"output": output_text}
     
     

requirements.txt

@@ -1,6 +1,6 @@
-fastapi>=0.74,<0.75
-requests>=2.27,<2.28
-uvicorn[standard]>=0.17,<0.18
+fastapi>=0.100.0
+requests>=2.31.0
+uvicorn[standard]>=0.24.0
 torch>=2.0
-transformers>=4.18
-sentencepiece==0.1.*
+transformers>=4.30.0
+sentencepiece>=0.2.0