Update app.py

app.py

@@ -5,6 +5,7 @@ from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel
 from huggingface_hub import InferenceClient
 import re
+from groq import Groq
 
 # Initialize FastAPI app
 app = FastAPI()
@@ -13,13 +14,16 @@ app = FastAPI()
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
 # Initialize Hugging Face Inference Client
-client = InferenceClient()
+#client = InferenceClient()
+
+client = Groq()
 
 # Pydantic model for API input
 class InfographicRequest(BaseModel):
     description: str
 
 # Load prompt template from environment variable
+SYSTEM_INSTRUCT = = os.getenv("SYSTEM_INSTRUCTOR")
 PROMPT_TEMPLATE = os.getenv("PROMPT_TEMPLATE")
 
 
@@ -41,6 +45,24 @@ async def extract_code_blocks(markdown_text):
     
     return code_blocks
 
+async def generate_infographic(infoRequest):
+    prompt = PROMPT_TEMPLATE.format(description=description)
+    generated_completion = client.chat.completions.create(
+            model="llama-3.1-70b-versatile",
+            messages=[
+                {"role": "system", "content": SYSTEM_INSTRUCT},
+                {"role": "user", "content": prompt}
+            ],
+            temperature=0.5,
+            max_tokens=5000,
+            top_p=1,
+            stream=False,
+            stop=None
+        )
+    generated_text = generated_completion.choices[0].message.content
+    code_blocks= await extract_code_blocks(generated_text)
+    return code_blocks
+
 # Route to serve the HTML template
 @app.get("/", response_class=HTMLResponse)
 async def serve_frontend():
@@ -53,23 +75,24 @@ async def generate_infographic(request: InfographicRequest):
     prompt = PROMPT_TEMPLATE.format(description=description)
 
     try:
-        messages = [{"role": "user", "content": prompt}]
-        stream = client.chat.completions.create(
-            model="Qwen/Qwen2.5-Coder-32B-Instruct",
-            messages=messages,
-            temperature=0.4,
-            max_tokens=6000,
-            top_p=0.7,
-            stream=True,
-        )
+        # messages = [{"role": "user", "content": prompt}]
+        # stream = client.chat.completions.create(
+        #     model="Qwen/Qwen2.5-Coder-32B-Instruct",
+        #     messages=messages,
+        #     temperature=0.4,
+        #     max_tokens=6000,
+        #     top_p=0.7,
+        #     stream=True,
+        # )
 
         
-        generated_text = ""
-        for chunk in stream:
-            generated_text += chunk.choices[0].delta.content
+        # generated_text = ""
+        # for chunk in stream:
+        #     generated_text += chunk.choices[0].delta.content
 
-        print(generated_text)
-        code_blocks= await extract_code_blocks(generated_text)
+        # print(generated_text)
+        #code_blocks= await extract_code_blocks(generated_text)
+        code_blocks=generate_infographic(description)
         if code_blocks:
             return JSONResponse(content={"html": code_blocks[0]})
         else: