Update inference

app.py

@@ -1,23 +1,17 @@
 import gradio as gr
-
 from langchain.embeddings import SentenceTransformerEmbeddings
 from langchain.vectorstores import Chroma
-
-from transformers import T5Tokenizer, T5ForConditionalGeneration
+from huggingface_hub import InferenceClient
 
 embeddings = SentenceTransformerEmbeddings(model_name="msmarco-distilbert-base-v4")
 db = Chroma(persist_directory="embeddings", embedding_function=embeddings)
 
-tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-large")
-model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large")
+client = InferenceClient(model="google/flan-t5-large")
 
 
 def respond(
         message,
         history: list[tuple[str, str]],
-        max_tokens,
-        temperature,
-        repetition_penalty,
 ):
     matching_docs = db.similarity_search(message)
 
@@ -37,28 +31,22 @@ def respond(
         f"Answer:"
     )
 
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
-
-    outputs = model.generate(input_ids,
-                             do_sample=True,
-                             max_new_tokens=max_tokens,
-                             temperature=temperature,
-                             repetition_penalty=repetition_penalty)
+    response = client.text_generation(
+        prompt,
+        max_new_tokens=250,
+        temperature=0.7,
+        top_p=0.95,
+    )
 
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    yield response
 
 
 demo = gr.ChatInterface(
     respond,
-    additional_inputs=[
-        gr.Slider(minimum=1, maximum=2048, value=1024, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(minimum=0.1, maximum=10, value=1.5, step=0.1, label="Repetition penalty"),
-    ],
     examples=[
-        "What types of roles are in the system?",
-        "How to import records into stock receipts in Boost.space?",
-        "Is it possible to create a PDF export from the product?",
+        ["What types of roles are in the system?"],
+        ["How to import records into stock receipts in Boost.space?"],
+        ["Is it possible to create a PDF export from the product?"],
     ],
 )