loads the lora

README.md

@@ -1,13 +1,15 @@
 ---
-title: Gpt Oss 20b Demo
-emoji: 💬
+title: GPT-OSS-20B Multilingual Reasoner Demo
+emoji: 🌟
 colorFrom: blue
 colorTo: pink
 sdk: gradio
 sdk_version: 5.40.0
 app_file: app.py
 pinned: false
-short_description: Try out OpenAI's gpt-oss-20b model
+short_description: Try out Tonic's GPT-OSS-20B Multilingual Reasoner LoRA adapter
 ---
 
+This demo showcases the GPT-OSS-20B model fine-tuned with LoRA for enhanced multilingual reasoning capabilities. The model is based on OpenAI's GPT-OSS-20B base model with a LoRA adapter from Tonic.
+
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,18 +1,69 @@
-from transformers import pipeline, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import torch
 from threading import Thread
 import gradio as gr
 import spaces
 import re
+from peft import PeftModel
 
-model_id = "Tonic/gpt-oss-20b-multilingual-reasoner"
+# Load the base model
+try:
+    base_model = AutoModelForCausalLM.from_pretrained(
+        "openai/gpt-oss-20b",
+        torch_dtype="auto",
+        device_map="auto",
+    )
+    tokenizer = AutoTokenizer.from_pretrained("openai/gpt-oss-20b")
+    
+    # Load the LoRA adapter
+    try:
+        model = PeftModel.from_pretrained(base_model, "Tonic/gpt-oss-20b-multilingual-reasoner")
+        print("✅ LoRA model loaded successfully!")
+    except Exception as lora_error:
+        print(f"⚠️ LoRA adapter failed to load: {lora_error}")
+        print("🔄 Falling back to base model...")
+        model = base_model
+        
+except Exception as e:
+    print(f"❌ Error loading model: {e}")
+    raise e
 
-pipe = pipeline(
-    "text-generation",
-    model=model_id,
-    torch_dtype="auto",
-    device_map="auto",
-)
+class LoRAPipeline:
+    def __init__(self, model, tokenizer):
+        self.model = model
+        self.tokenizer = tokenizer
+    
+    def __call__(self, messages, **kwargs):
+        prompt = self.format_messages(messages)
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+        
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                **kwargs
+            )
+        
+        generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        response = generated_text[len(prompt):]
+        return response
+    
+    def format_messages(self, messages):
+        """Format messages into a prompt string"""
+        formatted = ""
+        for message in messages:
+            role = message["role"]
+            content = message["content"]
+            if role == "system":
+                formatted += f"System: {content}\n"
+            elif role == "user":
+                formatted += f"User: {content}\n"
+            elif role == "assistant":
+                formatted += f"Assistant: {content}\n"
+        formatted += "Assistant: "
+        return formatted
+
+# Create the pipeline
+pipe = LoRAPipeline(model, tokenizer)
 
 def format_conversation_history(chat_history):
     messages = []
@@ -31,7 +82,7 @@ def generate_response(input_data, chat_history, max_new_tokens, system_prompt, t
     processed_history = format_conversation_history(chat_history)
     messages = system_message + processed_history + [new_message]
     
-    streamer = TextIteratorStreamer(pipe.tokenizer, skip_prompt=True, skip_special_tokens=True)
+    # Generate response using the LoRA pipeline
     generation_kwargs = {
         "max_new_tokens": max_new_tokens,
         "do_sample": True,
@@ -39,28 +90,48 @@ def generate_response(input_data, chat_history, max_new_tokens, system_prompt, t
         "top_p": top_p,
         "top_k": top_k,
         "repetition_penalty": repetition_penalty,
-        "streamer": streamer
+        "pad_token_id": tokenizer.eos_token_id,
     }
-    thread = Thread(target=pipe, args=(messages,), kwargs=generation_kwargs)
-    thread.start()
-    #streaming try #1
-    buffer = ""
+    
+    # For streaming, we'll generate token by token
+    prompt = pipe.format_messages(messages)
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    
+    # Generate with streaming
     full_response = ""
-    for chunk in streamer:
-        buffer += chunk
-        parts = re.split(r'(\s+)', buffer)
-        if re.match(r'\s+', parts[-1]) is not None:
-            to_append = ''.join(parts)
-            buffer = ""
-        else:
-            to_append = ''.join(parts[:-1])
-            buffer = parts[-1]
-        if to_append:
-            full_response += to_append
-            yield full_response
-    if buffer:
-        full_response += buffer
-        yield full_response
+    current_length = inputs["input_ids"].shape[1]
+    
+    with torch.no_grad():
+        for i in range(max_new_tokens):
+            # Generate one token at a time
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=1,
+                do_sample=True,
+                temperature=temperature,
+                top_p=top_p,
+                top_k=top_k,
+                repetition_penalty=repetition_penalty,
+                pad_token_id=tokenizer.eos_token_id,
+                use_cache=True
+            )
+            
+            # Get the new token
+            new_token = outputs[0][-1].unsqueeze(0)
+            
+            # Decode the new token
+            new_text = tokenizer.decode(new_token, skip_special_tokens=True)
+            
+            if new_text:
+                full_response += new_text
+                yield full_response
+            
+            # Update inputs for next iteration
+            inputs = {"input_ids": torch.cat([inputs["input_ids"], new_token], dim=1)}
+            
+            # Check for end of generation
+            if new_token.item() == tokenizer.eos_token_id:
+                break
 
 demo = gr.ChatInterface(
     fn=generate_response,
@@ -85,7 +156,7 @@ demo = gr.ChatInterface(
     cache_examples=False,
     type="messages",
     description="""
-# gpt-oss-20b
+# 🙋🏻‍♂️Welcome to 🌟Tonic's gpt-oss-20b Multilingual Reasoner Demo !
 Wait couple of seconds initially. You can adjust reasoning level in the system prompt like "Reasoning: high.
     """,
     fill_height=True,

requirements.txt

@@ -1,2 +1,4 @@
 git+https://github.com/huggingface/transformers.git
-accelerate
+accelerate
+peft
+torch