Update main.py

main.py

@@ -1,34 +1,16 @@
-from pydantic import BaseModel
+import requests
 
-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
-
-class Generate(BaseModel):
-    text: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)
+from fastapi import FastAPI
+
+API_URL = "https://api-inference.huggingface.co/models/google/gemma-1.1-7b-it"
+
+def query(payload):
+	response = requests.post(API_URL, json=payload)
+	return response.json()
+
+app = FastAPI()
 
 app.add_middleware(
     CORSMiddleware,
@@ -37,11 +19,12 @@ app.add_middleware(
     allow_methods=["*"],
     allow_headers=["*"],
 )
+    
 
 @app.get("/", tags=["Home"])
 def api_home():
     return {'detail': 'Welcome to FastAPI TextGen Tutorial!'}
 
-@app.post("/api/generate", summary="Generate text from prompt", tags=["Generate"], response_model=Generate)
-def inference(input_prompt: str):
-    return generate_text(prompt=input_prompt)
+@app.get("/api/generate", summary="Generate text from prompt", tags=["Generate"])
+def inference(input: str):
+    return query({"inputs": input})