Update main.py

main.py

@@ -1,23 +1,47 @@
-from transformers import pipeline, set_seed
-from fastapi import FastAPI
 from pydantic import BaseModel
-import uvicorn
 
+from .ConfigEnv import config
+from fastapi.middleware.cors import CORSMiddleware
 
+from langchain.llms import Clarifai
+from langchain.chains import LLMChain
+from langchain.prompts import PromptTemplate
+from TextGen import app
 
-app = FastAPI()
+class Generate(BaseModel):
+    text:str
 
-class Item(BaseModel):
-    prompt: str
-    
+def generate_text(prompt: str):
+    if prompt == "":
+        return {"detail": "Please provide a prompt."}
+    else:
+        prompt = PromptTemplate(template=prompt, input_variables=['Prompt'])
+        llm = Clarifai(
+            pat = config.CLARIFAI_PAT,
+            user_id = config.USER_ID,
+            app_id = config.APP_ID, 
+            model_id = config.MODEL_ID,
+            model_version_id=config.MODEL_VERSION_ID,
+        )
+        llmchain = LLMChain(
+            prompt=prompt,
+            llm=llm
+        )
+        llm_response = llmchain.run({"Prompt": prompt})
+        return Generate(text=llm_response)
 
-generator = pipeline('text-generation', model='gpt2')
-set_seed(42)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
-def generate(item: Item):
-    generator(item.prompt, max_length=30, num_return_sequences=5)
-   
 
-@app.post("/generate/")
-async def generate_text(item: Item):
-    return {"response": generate(item)}
+def api_home():
+    return {'detail': 'Welcome to FastAPI TextGen Tutorial!'}
+
+
+def inference(input_prompt: str):
+    return generate_text(prompt=input_prompt)