runai.py

import os
import json
import torch
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
)

MODEL_NAME = input("Enter model path or Hugging Face model name: ").strip()

print(f"\nLoading: {MODEL_NAME}")

model_type = "unknown"

# Try reading config.json manually
config_path = os.path.join(MODEL_NAME, "config.json")

if os.path.exists(config_path):
    try:
        with open(config_path, "r") as f:
            config_data = json.load(f)

        model_type = config_data.get("model_type", "unknown")
    except Exception:
        pass

# Fallback detection from path name
lower_name = MODEL_NAME.lower()

if model_type == "unknown":
    if "gpt2" in lower_name:
        model_type = "gpt2"
    elif "llama" in lower_name or "smollm" in lower_name:
        model_type = "llama"

print(f"Detected model type: {model_type}")

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

device = "cuda" if torch.cuda.is_available() else "cpu"

# Load model
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto" if torch.cuda.is_available() else None,
)

if not torch.cuda.is_available():
    model = model.to(device)

model.eval()

print("\nType 'exit' to quit.\n")

while True:
    user_input = input("You: ")

    if user_input.lower() == "exit":
        break

    inputs = tokenizer(
        user_input,
        return_tensors="pt"
    ).to(device)

    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=100,
            do_sample=True,
            temperature=0.8,
            top_p=0.95,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.eos_token_id,
        )

    response = tokenizer.decode(
        output[0],
        skip_special_tokens=True
    )

    if response.startswith(user_input):
        response = response[len(user_input):].strip()

    print(f"AI: {response}\n")