to cpu/ using transformers

inference.py

@@ -1,14 +1,14 @@
-from transformers import TextStreamer
+from transformers import TextStreamer, AutoModelForCausalLM, AutoTokenizer
 from typing import Tuple, List, Dict
 import torch
-from unsloth import FastLanguageModel
+# from unsloth import FastLanguageModel
 
 def load_model(
     model_name: str,
     max_seq_length: int,
     dtype: torch.dtype,
     load_in_4bit: bool
-) -> Tuple[FastLanguageModel, any]:
+) -> Tuple[AutoModelForCausalLM, any]:
     """
     Load and initialize the language model for inference.
     
@@ -21,13 +21,18 @@ def load_model(
     Returns:
         Tuple[FastLanguageModel, any]: Tuple containing the model and tokenizer
     """
-    model, tokenizer = FastLanguageModel.from_pretrained(
+    model_name = "lora_model"
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    model = AutoModelForCausalLM.from_pretrained(
         model_name=model_name,
-        max_seq_length=max_seq_length,
-        dtype=dtype,
-        load_in_4bit=load_in_4bit,
+        torch_dtype=dtype,
+        device_map="auto"
     )
-    FastLanguageModel.for_inference(model)
+
+    model.eval() # Set model to evaluation mode
+
     return model, tokenizer
 
 def prepare_input(
@@ -54,7 +59,7 @@ def prepare_input(
     ).to(device)
 
 def generate_response(
-    model: FastLanguageModel,
+    model: AutoModelForCausalLM,
     inputs: torch.Tensor,
     tokenizer: any,
     max_new_tokens: int = 2000,
@@ -77,16 +82,19 @@ def generate_response(
     Returns:
         str: Generated response
     """
-    text_streamer = TextStreamer(tokenizer, skip_prompt=skip_prompt)
+    # text_streamer = TextStreamer(tokenizer, skip_prompt=skip_prompt)
+    inputs = tokenizer(inputs, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
     outputs = model.generate(
-        input_ids=inputs,
-        streamer=text_streamer,
-        max_new_tokens=max_new_tokens,
-        use_cache=True,
-        temperature=temperature,
-        min_p=min_p
+        **inputs,
+        max_length=2000
+        # streamer=text_streamer,
+        # max_new_tokens=max_new_tokens,
+        # use_cache=True,
+        # temperature=temperature,
+        # min_p=min_p
     )
-    return outputs
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return generated_text
 
 def main(
         USER_INPUT_CODE = "program sum_of_numbers\n    implicit none\n    integer :: n, i, sum\n\n    ! Initialize variables\n    sum = 0\n\n    ! Get user input\n    print *, \"Enter a positive integer:\"\n    read *, n\n\n    ! Calculate the sum of numbers from 1 to n\n    do i = 1, n\n        sum = sum + i\n    end do\n\n    ! Print the result\n    print *, \"The sum of numbers from 1 to\", n, \"is\", sum\nend program sum_of_numbers",