.

app.py

@@ -16,18 +16,21 @@ peft_config = PeftConfig.from_pretrained("pseudolab/K23_MiniMed")
 peft_model = MistralForCausalLM.from_pretrained("pseudolab/K23_MiniMed", trust_remote_code=True)
 peft_model = PeftModel.from_pretrained(peft_model, "pseudolab/K23_MiniMed")
 
+text_generator = pipeline('text-generation', model=peft_model, tokenizer=tokenizer)
+
 # Prepare the context
 def prepare_context(data):
     # Format the data as a string
     data_str = data.to_string(index=False, header=False)
 
     # Tokenize the data
-    input_ids = tokenizer.encode(data_str, return_tensors="pt")
+    # input_ids = tokenizer.encode(data_str, return_tensors="pt")
 
     # Truncate the input if it's too long for the model
-    max_length = tokenizer.model_max_length
-    if input_ids.shape[1] > max_length:
-        input_ids = input_ids[:, :max_length]
+    # max_length = tokenizer.model_max_length
+    # if input_ids.shape[1] > max_length:
+    #     input_ids = input_ids[:, :max_length]
+    input_ids = data_str
 
     return input_ids
 
@@ -37,7 +40,8 @@ def fn(uploaded_file) -> str:
 
     # Generate text based on the context
     context = prepare_context(data)
-    generated_text = pipeline('text-generation', model=peft_model)(context)[0]['generated_text']
+    # generated_text = pipeline('text-generation', model=peft_model, tokenizer=tokenizer)(context)[0]['generated_text']
+    generated_text = text_generator(context)[0]['generated_text']
     ret += generated_text
 
     # Internally prompt the model to data analyze the EHR patient data
@@ -48,13 +52,14 @@ def fn(uploaded_file) -> str:
         input_ids = tokenizer.encode(prompt, return_tensors="pt")
 
         # Generate text based on the prompt
-        generated_text = pipeline('text-generation', model=peft_model)(input_ids=input_ids)[0]['generated_text']
+        # generated_text = pipeline('text-generation', model=peft_model, tokenizer=tokenizer)(input_ids=input_ids)[0]['generated_text']
+        generated_text = text_generator(prompt)[0]['generated_text']
         ret += generated_text
 
     return ret
 
 
-demo = gr.Interface(fn=fn, inputs="file", outputs="text")
+demo = gr.Interface(fn=fn, inputs="file", outputs="text", theme="pseudolab/huggingface-korea-theme")
 
 
 if __name__ == "__main__":