Revert "Replace use of transformers library with vLLM for high-speed inference"

This reverts commit 56c4132b099df59a010a2b8a2bc9ed88150da502.

quiz_generator.py

@@ -1,6 +1,6 @@
 import json
 import spaces
-from vllm import LLM, SamplingParams
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
 example_quiz = """
         {
@@ -32,14 +32,18 @@ system_prompt = f"""
         Final Answer:{example_quiz}
     """
 
-# Initialize vLLM model
+# Initialize model and pipeline
 model_id = "openai/gpt-oss-20b"
-llm = LLM(
-    model=model_id,
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
     dtype="auto",
-    tensor_parallel_size=1,  # Adjust based on your GPU setup
+    device_map="auto",
 )
 
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+
 @spaces.GPU(duration=90)
 def run_inference(prompt_messages):
     """
@@ -49,25 +53,14 @@ def run_inference(prompt_messages):
     See https://huggingface.co/docs/hub/en/spaces-zerogpu
 
     :param prompt_messages: The system and user messages submitted to the LLM
-    :return: Generated text from vLLM
+    :return: All messages returned by the LLM
     """
-    # Use the model's chat template to format messages
-    prompt = llm.get_tokenizer().apply_chat_template(
+    return pipe(
         prompt_messages,
-        tokenize=False,
-        add_generation_prompt=True
-    )
-    
-    # Set up sampling parameters
-    sampling_params = SamplingParams(
-        max_tokens=3000,
+        max_new_tokens=3000,
         temperature=0.7,
-        top_p=0.9,
+        do_sample=True,
     )
-    
-    # Generate response
-    outputs = llm.generate([prompt], sampling_params)
-    return outputs[0].outputs[0].text
 
 def to_final_answer(response):
     """
@@ -83,13 +76,17 @@ def to_final_answer(response):
     """
     first_json_key = '"questions":'
 
-    print('all_generated:', response)
-    last_marker_idx = response.rfind(first_json_key)
+    # Code from https://huggingface.co/docs/transformers/en/conversations#textgenerationpipeline
+    # The assistant response is always the last in the generated_text array, so -1.
+    assistant_response = response[0]["generated_text"][-1]["content"]
+
+    print('all_generated:', assistant_response)
+    last_marker_idx = assistant_response.rfind(first_json_key)
     if last_marker_idx != -1:
-        text = "{" + response[last_marker_idx:].strip()
+        text = "{" + assistant_response[last_marker_idx:].strip()
     else:
-        # Fallback: use the entire response
-        text = response.strip()
+        # Fallback: use the last response's text
+        text = response[-1]["generated_text"].strip()
     print('final text:', text)
     return text
 

requirements.txt

@@ -1,5 +1,5 @@
 gradio
 spaces
-vllm
+transformers
 torch
 accelerate