Update app.py

app.py

@@ -1,5 +1,4 @@
 import os
-import openai
 from openai import OpenAI
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from datasets import load_dataset, Dataset
@@ -11,6 +10,7 @@ from typing import List
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import ChatPromptTemplate
 import gradio as gr
+from huggingface_hub import InferenceClient
 
 # Configuration
 
@@ -25,7 +25,6 @@ DEFAULT_QUESTION = "Ask me anything in the context of persona-driven prompt gene
 os.environ['OPENAI_BASE'] = "https://api.openai.com/v1"
 os.environ['OPENAI_MODEL'] = "gpt-4"
 os.environ['MODEL_PROVIDER'] = "huggingface"
-api_key = os.environ.get("OPENAI_API_KEY")
 model_provider = os.environ.get("MODEL_PROVIDER")
 
 # Instantiate the client for openai v1.x
@@ -37,23 +36,10 @@ if model_provider.lower() == "openai":
     )
 else:
     MODEL_NAME = "meta-llama/Llama-3.3-70B-Instruct"
-    # Initialize Hugging Face client with streaming support
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=os.environ.get("HF_TOKEN"))
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME,
-        device_map='auto',
-        token=os.environ.get("HF_TOKEN"),
-        torch_dtype=torch.bfloat16,
-    )
-    streamer = TextStreamer(tokenizer, skip_prompt=True)
-    question_answerer = pipeline(
-        "text-generation",
-        model=model,
-        tokenizer=tokenizer,
-        device_map='auto',
-        streamer=streamer,
-        max_new_tokens=512,
-        return_full_text=False
+    # Initialize Hugging Face InferenceClient
+    hf_client = InferenceClient(
+        model=MODEL_NAME,
+        api_key=os.environ.get("HF_TOKEN")
     )
 
 # Load the Hugging Face dataset
@@ -96,18 +82,26 @@ def generate_response(question, history):
         print(f"[LOG] Generated prompt: {prompt[:200]}...")  # Log first 200 chars of prompt
         
         if model_provider.lower() == "huggingface":
-            prompt_template = """
-            <s>[INST] <<SYS>>
-            You are a helpful AI assistant. Answer the question based on the provided context.
-            <</SYS>>
+            messages = [
+                {
+                    "role": "system",
+                    "content": "You are a helpful AI assistant. Answer the question based on the provided context."
+                },
+                {
+                    "role": "user",
+                    "content": prompt
+                }
+            ]
+            
+            completion = hf_client.chat.completions.create(
+                model=MODEL_NAME,
+                messages=messages,
+                max_tokens=500
+            )
+            response = completion.choices[0].message.content
+            print(f"[LOG] Using Hugging Face model (serverless): {MODEL_NAME}")
+            print(f"[LOG] Hugging Face response: {response[:200]}...")
         
-            {prompt}[/INST]
-            """
-            chat_prompt = ChatPromptTemplate.from_template(prompt_template)
-            result = question_answerer(chat_prompt.format(prompt=prompt))
-            response = result[0]['generated_text'] if isinstance(result, list) else result
-            print(f"[LOG] Using Hugging Face model: {MODEL_NAME}")
-            print(f"[LOG] Hugging Face response: {response[:200]}...")  # Log first 200 chars of response
         elif model_provider.lower() == "openai":
             response = client.chat.completions.create(
                 model=os.environ.get("OPENAI_MODEL"),