Update app.py

app.py

@@ -1,38 +1,27 @@
 import gradio as gr
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
-from threading import Thread
 import os
 import requests
 
-
-class StopOnTokens(StoppingCriteria):
-    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
-        stop_ids = [29, 0]
-        for stop_id in stop_ids:
-            if input_ids[0][-1] == stop_id:
-                return True
-        return False
+# Define the Hugging Face API details
+API_URL = "https://api-inference.huggingface.co/models/Huzaifa367/chat-summarizer"
+API_TOKEN = os.getenv("AUTH_TOKEN")
+HEADERS = {"Authorization": f"Bearer {API_TOKEN}"}
 
 def predict(message, history):
+    # Prepare input data for the Hugging Face API
     history_transformer_format = history + [[message, ""]]
-    stop = StopOnTokens()
-
-    messages = "".join(["".join(["\n<human>:"+item[0], "\n<bot>:"+item[1]])
-                for item in history_transformer_format])
+    messages = "".join([f"<human>:{item[0]}\n<bot>:{item[1]}" for item in history_transformer_format])
 
-    # Use your API to generate the response
     try:
-        API_URL = "https://api-inference.huggingface.co/models/Huzaifa367/chat-summarizer"
-        API_TOKEN = os.getenv("AUTH_TOKEN")
-        HEADERS = {"Authorization": f"Bearer {API_TOKEN}"}
+        # Send request to the Hugging Face API
         payload = {"inputs": messages}
-        
         response = requests.post(API_URL, headers=HEADERS, json=payload)
-        response.raise_for_status()  # Raise exception for non-2xx status codes       
-        summary_text = response.json()[0]["summary_text"] if isinstance(response, list) else response.get("summary_text", "")
+        response.raise_for_status()  # Raise exception for non-2xx status codes
+
+        # Extract summary text from the response
+        summary_text = response.json()["summary_text"]
     except requests.exceptions.RequestException as e:
-        print(f"Error querying Hugging Face API: {e}")
+        # Handle API request errors
         summary_text = f"Error querying Hugging Face API: {e}"
     
     return summary_text