Update app.py

app.py

@@ -2,29 +2,38 @@ from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipe
 import torch
 import pickle
 import streamlit as st
+from translate import Translator
+
 from huggingface_hub import InferenceClient
+import gradio as gr
 
-client = InferenceClient(
-    "mistralai/Mistral-7B-Instruct-v0.1"
-)
 
 
-device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
 
-from translate import Translator
 
-def init_session_state():
-    if 'history' not in st.session_state:
-        st.session_state.history = ""
 
+def classify(text,labels):
+    output = classifier(text, labels, multi_label=False)
+    
+    return output
 
-temperature=0.9
-max_new_tokens=256
-top_p=0.95
-repetition_penalty=1.0
+def format_prompt(message, history):
+  prompt = "<s>"
+  for user_prompt, bot_response in history:
+    prompt += f"[INST] {user_prompt} [/INST]"
+    prompt += f" {bot_response}</s> "
+  prompt += f"[INST] {message} [/INST]"
+  return prompt
 
+def generate(
+    prompt, history, temperature=0.7, max_new_tokens=256, top_p=0.95, repetition_penalty=1.1,
+):
+    temperature = float(temperature)
+    if temperature < 1e-2:
+        temperature = 1e-2
+    top_p = float(top_p)
 
-generate_kwargs = dict(
+    generate_kwargs = dict(
         temperature=temperature,
         max_new_tokens=max_new_tokens,
         top_p=top_p,
@@ -33,89 +42,83 @@ generate_kwargs = dict(
         seed=42,
     )
 
-def format_prompt(message, history):
-  prompt = "<s>"
-  for user_prompt, bot_response in history:
-    prompt += f"[INST] {user_prompt} [/INST]"
-    prompt += f" {bot_response}</s> "
-  prompt += f"[INST] {message} [/INST]"
-  return prompt
-
+    formatted_prompt = format_prompt(prompt, history)
 
-# Initialize session state
-init_session_state()
+    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
+    output = ""
 
-# pipe = pipeline("text2text-generation", model="google/flan-t5-base")
-pipe = pipeline("text-generation", model="microsoft/Phi-3-mini-4k-instruct", trust_remote_code=True)
-# pipe = pipeline("text-generation", model="GeneZC/MiniChat-1.5-3B")
-# pipe = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.2")
-# model_name = "MoritzLaurer/mDeBERTa-v3-base-mnli-xnli"
-# tokenizer = AutoTokenizer.from_pretrained(model_name)
-# model = AutoModelForSequenceClassification.from_pretrained(model_name)
+    for response in stream:
+        output += response.token.text
+        yield output
+    return output
 
-classifier = pipeline("zero-shot-classification", model="MoritzLaurer/mDeBERTa-v3-base-mnli-xnli")
 
-# with open('chapter_titles.pkl', 'rb') as file:
-#     titles_astiko = pickle.load(file)
-# labels1 = ["κληρονομικό", "ακίνητα", "διαζύγιο"]
-# # labels2 = ["αποδοχή κληρονομιάς", "αποποίηση", "διαθήκη"]
-# # labels3 = ["μίσθωση", "κυριότητα", "έξωση", "απλήρωτα νοίκια"]
+client = InferenceClient(
+    "mistralai/Mistral-7B-Instruct-v0.1"
+)
 
+classifier = pipeline("zero-shot-classification", model="MoritzLaurer/mDeBERTa-v3-base-mnli-xnli")
 
-# titles_astiko = ["γάμος", "αλλοδαπός", "φορολογία", "κληρονομικά", "στέγη", "οικογενειακό", "εμπορικό","κλοπή","απάτη"]
-# Load dictionary from the file using pickle
 with open('my_dict.pickle', 'rb') as file:
     dictionary = pickle.load(file)
 
-def classify(text,labels):
-    output = classifier(text, labels, multi_label=False)
     
-    return output
-
+# text = st.text_input('Enter some text:')  # Input field for new text
 
-text = st.text_input('Enter some text:')  # Input field for new text
 
-if text:
-
-    labels = list(dictionary)
+#     labels = list(dictionary)
     
-    output = classify(text,labels)
+#     output = classify(text,labels)
 
-    output = output["labels"][0]
+#     output = output["labels"][0]
 
-    labels = list(dictionary[output])
+#     labels = list(dictionary[output])
 
-    output2 = classify(text,labels)
+#     output2 = classify(text,labels)
 
-    output2 = output2["labels"][0]
+#     output2 = output2["labels"][0]
 
 
-    answer = dictionary[output][output2]
+#     answer = dictionary[output][output2]
 
-    # Create a translator object with specified source and target languages
-    translator = Translator(from_lang='el', to_lang='en')
-    translator2 = Translator(from_lang='en', to_lang='el')
+#     # Create a translator object with specified source and target languages
+#     translator = Translator(from_lang='el', to_lang='en')
+#     translator2 = Translator(from_lang='en', to_lang='el')
 
-    st.text("H ερώτηση σας σχετίζεται με " + output+ " δίκαιο")
+#     st.text("H ερώτηση σας σχετίζεται με " + output+ " δίκαιο")
  
 
-# Translate the text from Greek to English
-    answer = translator.translate(answer)
-    text = translator.translate(text)
+# # Translate the text from Greek to English
+#     answer = translator.translate(answer)
+#     text = translator.translate(text)
 
-    st.text("Πιο συγκεκριμένα σχετίζεται με " + output2)
+#     st.text("Πιο συγκεκριμένα σχετίζεται με " + output2)
+    
 
+    
 
-# text_to_translate2 = text[499:999]
-# translated_text2 = translator.translate(text_to_translate2)
+#     prompt =  "Based on this info only:" + answer +" ,answer this question, by reasoning step by step:" + text
+#     formatted_prompt = format_prompt(prompt, history)
+    # translated_text2 = translator2.translate(output)
 
-    
 
-    # st.session_state.history += "Based on this info only:" + answer +" ,answer this question, by reasoning step by step:" + text  # Add new text to history
-    # out = pipe(st.session_state.history, max_new_tokens=256)  # Generate output based on history
+def generate(
+    prompt, history, temperature=0.7, max_new_tokens=256, top_p=0.95, repetition_penalty=1.1,
+):
+    temperature = float(temperature)
+    if temperature < 1e-2:
+        temperature = 1e-2
+    top_p = float(top_p)
+
+    generate_kwargs = dict(
+        temperature=temperature,
+        max_new_tokens=max_new_tokens,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty,
+        do_sample=True,
+        seed=42,
+    )
 
-    history = st.session_state.history
-    prompt =  "Based on this info only:" + answer +" ,answer this question, by reasoning step by step:" + text
     formatted_prompt = format_prompt(prompt, history)
 
     stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
@@ -125,22 +128,69 @@ if text:
         output += response.token.text
         yield output
     return output
+    
 
 
-    # st.text(st.session_state.history)
-    
-    # translated_text2 = translator2.translate(out[0]['generated_text'])
-    translated_text2 = translator2.translate(output)
 
-    st.text(translated_text2)
 
-    # with st.expander("View Full Output", expanded=False):
-    #     st.write(translated_text2, allow_output_mutation=True)
 
-    # st.text(translated_text2)
-    # st.text("History: " + st.session_state.history)
 
-    # st.text(output)
-    # st.text(output2)
 
-    # st.text(answer)
+additional_inputs=[
+    gr.Slider(
+        label="Temperature",
+        value=0.7,
+        minimum=0.0,
+        maximum=1.0,
+        step=0.05,
+        interactive=True,
+        info="Higher values produce more diverse outputs",
+    ),
+    gr.Slider(
+        label="Max new tokens",
+        value=256,
+        minimum=0,
+        maximum=1024,
+        step=64,
+        interactive=True,
+        info="The maximum numbers of new tokens",
+    ),
+    gr.Slider(
+        label="Top-p (nucleus sampling)",
+        value=0.95,
+        minimum=0.0,
+        maximum=1,
+        step=0.05,
+        interactive=True,
+        info="Higher values sample more low-probability tokens",
+    ),
+    gr.Slider(
+        label="Repetition penalty",
+        value=1.1,
+        minimum=1.0,
+        maximum=2.0,
+        step=0.05,
+        interactive=True,
+        info="Penalize repeated tokens",
+    )
+]
+
+css = """
+  #mkd {
+    height: 500px; 
+    overflow: auto; 
+    border: 1px solid #ccc; 
+  }
+"""
+
+with gr.Blocks(css=css) as demo:
+    gr.HTML("<h1><center>Mistral 7B Instruct<h1><center>")
+    gr.HTML("<h3><center>In this demo, you can chat with <a href='https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1'>Mistral-7B-Instruct</a> model. 💬<h3><center>")
+    gr.HTML("<h3><center>Learn more about the model <a href='https://huggingface.co/docs/transformers/main/model_doc/mistral'>here</a>. 📚<h3><center>")
+    gr.ChatInterface(
+        generate,
+        additional_inputs=additional_inputs,
+        examples=[["What is the secret to life?"], ["Write me a recipe for pancakes."]]
+    )
+
+demo.queue().launch(debug=True)