Using MistralAI API

app.py

@@ -6,18 +6,29 @@ Description: Using a large language model for translation.
 Author: Didier Guillevic
 Date: 2024-09-17
 """
-import spaces
-import torch
-import gradio as gr
-import transformers
-from transformers import TextIteratorStreamer
-from threading import Thread
-
 import logging
 logger = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO)
 
-from model_llm import tokenizer, model, model_name
+import gradio as gr
+
+#import spaces
+#import torch
+#import transformers
+#from transformers import TextIteratorStreamer
+#from threading import Thread
+#from model_llm import tokenizer, model
+
+import os
+from mistralai import Mistral
+
+#
+# Mistral AI client
+#
+api_key = os.environ["MISTRAL_API_KEY"]
+client = Mistral(api_key=api_key)
+model_id = "mistral-large-latest" # 128k context window
+
 
 #
 # Default instruction: translate given text
@@ -30,9 +41,8 @@ translation_instruction = (
 )
 
 #
-# Generate a response given some user input and optional chat history
+# Generate a response using MistralAI API
 #
-@spaces.GPU
 def generate_chat_response_streaming(
         input_text,
         chat_history=None,
@@ -41,6 +51,40 @@ def generate_chat_response_streaming(
         temperature=0.0
     ):
     """Given some input from the user (and a chat history), generate a response"""
+
+    # messages up to now
+    messages = []
+    if not chat_history:
+        messages.append({'role': 'user', 'content': instruction_message + ' ' + input_text})
+    else:
+        for input, response in chat_history:
+            messages.append({"role": "user", "content": input})
+            messages.append({"role": "assistant", "content": response})
+        messages.append({'role': 'user', 'content': input_text})
+    logger.info(messages)
+
+    # generate response
+    
+    # Yield the model response as the tokens are being generated
+    stream_reponse = client.chat.stream(model=model_id, messages=messages)
+    model_response = ""
+    for chunk in stream_reponse:
+        model_response += chunk.data.choices[0].delta.content
+        yield model_response
+
+
+#
+# Generate a response given some user input and optional chat history
+#
+#@spaces.GPU
+def generate_chat_response_streaming__(
+        input_text,
+        chat_history=None,
+        instruction_message=translation_instruction,
+        max_new_tokens=1_024,
+        temperature=0.0
+    ):
+    """Given some input from the user (and a chat history), generate a response"""
     
     # Conversation up to now
     conversation = []
@@ -136,7 +180,7 @@ with gr.Blocks() as demo:
         textbox=textbox,
         #clear_btn=None,  # Unfortunately, clear_btn also reset the additional inputs. Hence disabling for now.
         examples=examples,
-        cache_examples=True,
+        cache_examples=False,
         retry_btn="Retry",
         undo_btn="Undo",
         clear_btn="Clear",
@@ -152,4 +196,4 @@ with gr.Blocks() as demo:
         """)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(show_api=False)