Update app.py

app.py

@@ -1,49 +1,48 @@
 import gradio as gr
-from huggingface_hub import snapshot_download
-from vllm import LLM, SamplingParams
+from datetime import datetime
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
 # ------------------------
 # 1) Load the Model
 # ------------------------
-# Download the model repository, specify revision if needed
-model_dir = snapshot_download(repo_id="BSC-LT/salamandraTA-7B-instruct-GGUF", revision="main",     allow_patterns=[
-                                                                                                                    "salamandrata_7b_inst_q4.gguf", 
-                                                                                                                    "*tokenizer*", 
-                                                                                                                    "tokenizer_config.json", 
-                                                                                                                    "tokenizer.model", 
-                                                                                                                    "config.json",
-                                                                                                                ])
-model_name = "salamandrata_7b_inst_q4.gguf"
-
-# Create an LLM instance from vLLM
-llm = LLM(model=model_dir + '/' + model_name, tokenizer=model_dir)
-
-# We can define a single helper function to call the model:
+model_id = "BSC-LT/salamandraTA-7b-instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    device_map="auto",
+    torch_dtype=torch.bfloat16
+)
+
+# Common function to generate text using transformers
 def call_model(prompt: str, temperature: float = 0.1, max_tokens: int = 256):
-    """
-    Sends the prompt to the LLM using vLLM's chat interface.
-    """
-    messages = [{'role': 'user', 'content': prompt}]
-    outputs = llm.chat(
-        messages,
-        sampling_params=SamplingParams(
-            temperature=temperature,
-            stop_token_ids=[5],  # you can adjust the stop token ID if needed
-            max_tokens=max_tokens
-        )
+    message = [{"role": "user", "content": prompt}]
+    date_string = datetime.today().strftime('%Y-%m-%d')
+
+    chat_prompt = tokenizer.apply_chat_template(
+        message,
+        tokenize=False,
+        add_generation_prompt=True,
+        date_string=date_string
+    )
+
+    inputs = tokenizer.encode(chat_prompt, return_tensors="pt").to(model.device)
+    input_length = inputs.shape[1]
+    outputs = model.generate(
+        input_ids=inputs,
+        max_new_tokens=max_tokens,
+        do_sample=True,
+        temperature=temperature,
+        num_beams=5,
+        early_stopping=True
     )
-    # The model returns a list of "Generation" objects, each containing .outputs
-    return outputs[0].outputs[0].text if outputs else ""
+    return tokenizer.decode(outputs[0, input_length:], skip_special_tokens=True)
 
 # ------------------------
 # 2) Task-specific functions
 # ------------------------
 
 def general_translation(source_lang, target_lang, text):
-    """
-    General translation prompt:
-    Translate from source_lang into target_lang.
-    """
     prompt = (
         f"Translate the following text from {source_lang} into {target_lang}.\n"
         f"{source_lang}: {text}\n"
@@ -52,10 +51,6 @@ def general_translation(source_lang, target_lang, text):
     return call_model(prompt, temperature=0.1)
 
 def post_editing(source_lang, target_lang, source_text, machine_translation):
-    """
-    Post-editing prompt:
-    Ask the model to fix any mistakes in the machine translation or keep it unedited.
-    """
     prompt = (
         f"Please fix any mistakes in the following {source_lang}-{target_lang} machine translation or keep it unedited if it's correct.\n"
         f"Source: {source_text}\n"
@@ -65,10 +60,6 @@ def post_editing(source_lang, target_lang, source_text, machine_translation):
     return call_model(prompt, temperature=0.1)
 
 def document_level_translation(source_lang, target_lang, document_text):
-    """
-    Document-level translation prompt:
-    Translate a multi-paragraph document.
-    """
     prompt = (
         f"Please translate this text from {source_lang} into {target_lang}.\n"
         f"{source_lang}: {document_text}\n"
@@ -77,16 +68,7 @@ def document_level_translation(source_lang, target_lang, document_text):
     return call_model(prompt, temperature=0.1)
 
 def named_entity_recognition(tokenized_text):
-    """
-    Named-entity recognition prompt:
-    Label tokens as ORG, PER, LOC, MISC, or O.
-    Expects the user to provide a list of tokens.
-    """
-    # Convert the input string into a list of tokens, if the user typed them as space-separated words
-    # or if the user provided them as a Python list string, we can try to parse that.
-    # For simplicity, let's assume it's a space-separated string. 
     tokens = tokenized_text.strip().split()
-    
     prompt = (
         "Analyse the following tokenized text and mark the tokens containing named entities.\n"
         "Use the following annotation guidelines with these tags for named entities:\n"
@@ -102,10 +84,6 @@ def named_entity_recognition(tokenized_text):
     return call_model(prompt, temperature=0.1)
 
 def grammar_checker(source_lang, sentence):
-    """
-    Grammar checker prompt:
-    Fix any mistakes in the given source_lang sentence or keep it unedited if correct.
-    """
     prompt = (
         f"Please fix any mistakes in the following {source_lang} sentence or keep it unedited if it's correct.\n"
         f"Sentence: {sentence}\n"