Update app.py

app.py

@@ -1,63 +1,44 @@
 import gradio as gr
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
-from threading import Thread
-import spaces
-
-finetuned_model = "CONCREE/adia-llm"
-
-# Charge le modele
-model = AutoModelForCausalLM.from_pretrained(
-        finetuned_model,
-        device_map="auto",
-        trust_remote_code=True,
+from openai import OpenAI
+
+BASE_URL = "https://kks679fhv1td67-8000.proxy.runpod.net/v1"
+API_KEY="SOMEHOW"
+
+# Create an OpenAI client to interact with the API server
+client = OpenAI(
+    base_url=BASE_URL,
+    api_key=API_KEY
 )
 
-# Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained(finetuned_model,
-                                         trust_remote_code=True,
-                                         padding=True,
-                                         truncation=True)
-
-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
-
-@spaces.GPU
+
 def predict(message, history):
-    history_transformer_format = history + [[message, ""]]
-    stop = StopOnTokens()
-
-    messages = "".join(["".join(["\n[INST]:"+item[0], "\n[/INST]:"+item[1]]) for item in history_transformer_format])
-
-    model_inputs = tokenizer([messages], return_tensors="pt").to("cuda")
-    streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
-    generate_kwargs = dict(
-        model_inputs,
-        streamer=streamer,
-        max_new_tokens=1024,
-        num_beams=1,
-        stopping_criteria=StoppingCriteriaList([stop])
+    # Convert chat history to OpenAI format
+    history_openai_format = [{
+        "role": "system",
+        "content": "Tu es un excellent assistant IA nommé Adia, développé par CONCREE pour accompagner les entrepreneurs Africains."
+    }]
+    for human, assistant in history:
+        history_openai_format.append({"role": "user", "content": human})
+        history_openai_format.append({
+            "role": "assistant",
+            "content": assistant
+        })
+    history_openai_format.append({"role": "user", "content": message})
+
+    # Create a chat completion request and send it to the API server
+    stream = client.chat.completions.create(
+        model="CONCREE/meta-adia-llm-instruct",  # Model name to use
+        messages=history_openai_format,  # Chat history
+        temperature=0.1,  # Temperature for text generation
+        stream=True,  # Stream response
     )
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
-    partial_message = ""
-    start_flag = True  # Flag to ignore initial newline
 
-    for new_token in streamer:
-        if start_flag and new_token == '\n':
-            continue
-        start_flag = False
-        partial_message += new_token
+    # Read and return generated text from response stream
+    partial_message = ""
+    for chunk in stream:
+        partial_message += (chunk.choices[0].delta.content or "")
         yield partial_message
 
 
-demo = gr.ChatInterface(predict).launch()
-
-
-if __name__ == "__main__":
-    demo.launch()
+# Create and launch a chat interface with Gradio
+gr.ChatInterface(predict).queue().launch()