Add paper config

app.py

@@ -1,7 +1,13 @@
 import gradio as gr
 
 from peft import PeftModel, PeftConfig
-from transformers import MistralForCausalLM, TextIteratorStreamer, AutoTokenizer, BitsAndBytesConfig
+from transformers import (
+    MistralForCausalLM,
+    TextIteratorStreamer,
+    AutoTokenizer,
+    BitsAndBytesConfig,
+    GenerationConfig,
+)
 from time import sleep
 from threading import Thread
 from torch import float16
@@ -15,38 +21,60 @@ quant_config = BitsAndBytesConfig(
     bnb_4bit_use_double_quant=False,
 )
 
-model = MistralForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1",
-    quantization_config=quant_config)
-    #device_map="auto",)
+model = MistralForCausalLM.from_pretrained(
+    "mistralai/Mistral-7B-v0.1", quantization_config=quant_config
+)
+
+
+# device_map="auto",)
 model = PeftModel.from_pretrained(model, "lang-uk/dragoman").to("cuda")
-tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1", use_fast=False, add_bos_token=False)
+tokenizer = AutoTokenizer.from_pretrained(
+    "mistralai/Mistral-7B-v0.1", use_fast=False, add_bos_token=False
+)
+
 
 @spaces.GPU(duration=30)
 def translate(input_text):
     generated_text = ""
     input_text = input_text.strip()
-    for chunk in input_text.split("\n"):
-        if not chunk:
-            generated_text += "\n"
-            yield generated_text
-            continue
-        chunk = f"[INST] {chunk} [/INST]"
-        inputs = tokenizer([chunk], return_tensors="pt").to(model.device)
-
-        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-        generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=200)
-
-        thread = Thread(target=model.generate, kwargs=generation_kwargs)
-
-        thread.start()
-        
-        for new_text in streamer:
-            generated_text += new_text
-            yield generated_text
-        
-        generated_text += "\n"
-        yield generated_text
-
-
-iface = gr.Interface(fn=translate, inputs="text", outputs="text", examples=[["who holds this neighborhood?"]])
-iface.launch()
+
+    input_text = f"[INST] {input_text} [/INST]"
+    inputs = tokenizer([input_text], return_tensors="pt").to(model.device)
+
+    generation_kwargs = dict(inputs, max_new_tokens=200, num_beams=10, temperature=1) # streamer=streamer,
+    
+    
+    # streaming support
+    #streamer = TextIteratorStreamer(
+    #    tokenizer, skip_prompt=True, skip_special_tokens=True
+    #)
+
+    #thread = Thread(target=model.generate, kwargs=generation_kwargs)
+
+    #thread.start()
+
+
+    #for new_text in streamer:
+    #    generated_text += new_text
+    #    yield generated_text
+
+    #generated_text += "\n"
+    #yield generated_text
+
+    output = model.generate(**generation_kwargs)
+    output = tokenizer.decode(output[0], skip_special_tokens=True).split("[/INST] ")[-1].strip()
+    return output
+
+
+iface = gr.Interface(
+    fn=translate,
+    inputs=gr.Textbox(
+        value="",
+        label="Source sentence",
+    ),
+    outputs=gr.Textbox(label="Translated sentence"),
+    examples=[[
+        "ChatGPT (Chat Generative Pre-trained Transformer) is a chatbot developed by OpenAI and launched on November 30, 2022. Based on a large language model, it enables users to refine and steer a conversation towards a desired length, format, style, level of detail, and language. Successive prompts and replies, known as prompt engineering, are considered at each conversation stage as a context.[2] ",
+        "who holds this neighborhood?"]],
+)
+iface.launch()